From 1a65828365fab827633fa07566b60c5deb589160 Mon Sep 17 00:00:00 2001 From: Mainder Singh Date: Fri, 3 Mar 2017 19:45:15 +0530 Subject: [PATCH] DATAMONGO-1637 - Add support for aggregation result streaming. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 results = mongoOperations.aggregateStream(aggregation, "inputCollection", TagCount.class); List tagCount = new ArrayList(); while (results.hasNext()) { tagCount.add(results.next()); } results.close(); Original pull request: #447. --- .../data/mongodb/core/MongoOperations.java | 26 ++- .../data/mongodb/core/MongoTemplate.java | 69 ++++++++ .../core/aggregation/AggregationTests.java | 150 +++++++++++++++++- 3 files changed, 242 insertions(+), 3 deletions(-) diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoOperations.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoOperations.java index fad14684b..5eef1b52c 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoOperations.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoOperations.java @@ -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 { */ AggregationResults aggregate(TypedAggregation aggregation, String collectionName, Class outputType); + CloseableIterator aggregateStream(TypedAggregation aggregation, String inputCollectionName, + Class 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 { */ AggregationResults aggregate(TypedAggregation aggregation, Class 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 + */ + CloseableIterator aggregateStream(TypedAggregation aggregation, Class outputType); + + /** * Execute an aggregation operation. The raw results will be mapped to the given entity class. * @@ -404,6 +421,8 @@ public interface MongoOperations { */ AggregationResults aggregate(Aggregation aggregation, Class inputType, Class outputType); + CloseableIterator aggregateStream(Aggregation aggregation, Class inputType, Class outputType); + /** * Execute an aggregation operation. The raw results will be mapped to the given entity class. * @@ -417,6 +436,9 @@ public interface MongoOperations { */ AggregationResults aggregate(Aggregation aggregation, String collectionName, Class outputType); + CloseableIterator aggregateStream(Aggregation aggregation, String collectionName, Class 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 findAndModify(Query query, Update update, FindAndModifyOptions options, Class 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 diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java index b1d32978f..3bbc60a41 100644 --- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java +++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java @@ -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 CloseableIterator aggregateStream(TypedAggregation aggregation, Class outputType) { + return aggregateStream(aggregation, determineCollectionName(aggregation.getInputType()), outputType); + } + @Override public AggregationResults aggregate(TypedAggregation aggregation, String inputCollectionName, Class outputType) { @@ -1540,6 +1548,17 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware, return aggregate(aggregation, inputCollectionName, outputType, context); } + @Override + public CloseableIterator aggregateStream(TypedAggregation aggregation, String inputCollectionName, + Class 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 AggregationResults aggregate(Aggregation aggregation, Class inputType, Class outputType) { @@ -1547,11 +1566,23 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware, new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper)); } + @Override + public CloseableIterator aggregateStream(Aggregation aggregation, Class inputType, Class outputType) { + + return aggregateStream(aggregation, determineCollectionName(inputType), outputType, + new TypeBasedAggregationOperationContext(inputType, mappingContext, queryMapper)); + } + @Override public AggregationResults aggregate(Aggregation aggregation, String collectionName, Class outputType) { return aggregate(aggregation, collectionName, outputType, null); } + @Override + public CloseableIterator aggregateStream(Aggregation aggregation, String collectionName, Class 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 CloseableIterator aggregateStream(final Aggregation aggregation, final String collectionName, + final Class 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>() { + + @Override + public CloseableIterator doInCollection(DBCollection collection) throws MongoException, DataAccessException { + + List pipeline = (List) command.get("pipeline"); + Cursor cursor = collection.aggregate(pipeline, getNativeAggregationOptionsFromCommand(command)); + + ReadDbObjectCallback readCallback = new ReadDbObjectCallback(mongoConverter, outputType, collectionName); + + return new CloseableIterableCursorAdapter(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; diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java index 85a38a5b3..1ec9d504a 100644 --- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java +++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java @@ -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 MongoDB Aggregation Examples */ 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 iterator = mongoTemplate.aggregateStream(agg, INPUT_COLLECTION, TagCount.class); + + assertThat(iterator, is(notNullValue())); + List tagCount = new ArrayList(); + 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 results = mongoTemplate.aggregateStream(aggregation, INPUT_COLLECTION, TagCount.class); + + assertThat(results, is(notNullValue())); + + List tagCount = new ArrayList(); + 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 results = mongoTemplate.aggregateStream(aggregation, INPUT_COLLECTION, TagCount.class); + + assertThat(results, is(notNullValue())); + + List tagCount = new ArrayList(); + 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 agg = createUsersWithCommonLikesAggregation() // + .withOptions(newAggregationOptions().allowDiskUse(true).build()); + + assertThat(agg, is(notNullValue())); + assertThat(agg.toString(), is(notNullValue())); + + CloseableIterator iterator = mongoTemplate.aggregateStream(agg, LikeStats.class); + List result = new ArrayList(); + 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 agg = newAggregation(Person.class, // + group("sex").count().as("count"), // + sort(DESC, "count"), // + out(tempOutCollection)); + + CloseableIterator iterator = mongoTemplate.aggregateStream(agg, DBObject.class); + + List 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; } } +