Allow to estimate document count.

This commit introduce an option that allows users to opt in on using estimatedDocumentCount instead of countDocuments in case the used filter query is empty.

Author: christophstrobl

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

@@ -22,13 +22,13 @@
 import java.math.RoundingMode;
 import java.util.*;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiPredicate;
 import java.util.stream.Collectors;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.bson.Document;
 import org.bson.conversions.Bson;
-
 import org.springframework.beans.BeansException;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
@@ -188,6 +188,11 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 
 	private SessionSynchronization sessionSynchronization = SessionSynchronization.ON_ACTUAL_TRANSACTION;
 
+	private CountExecution countExecution = (collectionName, filter, options) -> {
+		return execute(collectionName,
+				collection -> collection.countDocuments(CountQuery.of(filter).toQueryDocument(), options));
+	};
+
 	/**
 	 * Constructor used for a basic template configuration.
 	 *
@@ -345,6 +350,51 @@ public void setEntityCallbacks(EntityCallbacks entityCallbacks) {
 		this.entityCallbacks = entityCallbacks;
 	}
 
+	/**
+	 * En-/Disable usage of estimated count.
+	 *
+	 * @param enabled if {@literal true} {@link MongoCollection#estimatedDocumentCount()} ()} will we used for unpaged,
+	 *          empty {@link Query queries}.
+	 * @since 3.4
+	 */
+	public void useEstimatedCount(boolean enabled) {
+		useEstimatedCount(enabled, this::countCanBeEstimated);
+	}
+
+	/**
+	 * En-/Disable usage of estimated count based on the given {@link BiPredicate estimationFilter}.
+	 *
+	 * @param enabled if {@literal true} {@link MongoCollection#estimatedDocumentCount()} will we used for {@link Document
+	 *          filter queries} that pass the given {@link BiPredicate estimationFilter}.
+	 * @param estimationFilter the {@link BiPredicate filter}.
+	 * @since 3.4
+	 */
+	public void useEstimatedCount(boolean enabled, BiPredicate<Document, CountOptions> estimationFilter) {
+
+		if (enabled) {
+
+			this.countExecution = (collectionName, filter, options) -> {
+
+				if (!estimationFilter.test(filter, options)) {
+					return execute(collectionName,
+							collection -> collection.countDocuments(CountQuery.of(filter).toQueryDocument(), options));
+				}
+
+				EstimatedDocumentCountOptions estimatedDocumentCountOptions = new EstimatedDocumentCountOptions();
+				if (options.getMaxTime(TimeUnit.MILLISECONDS) > 0) {
+					estimatedDocumentCountOptions.maxTime(options.getMaxTime(TimeUnit.MILLISECONDS), TimeUnit.MILLISECONDS);
+				}
+
+				return doEstimatedCount(collectionName, estimatedDocumentCountOptions);
+			};
+		} else {
+			this.countExecution = (collectionName, filter, options) -> {
+				return execute(collectionName,
+						collection -> collection.countDocuments(CountQuery.of(filter).toQueryDocument(), options));
+			};
+		}
+	}
+
 	/**
 	 * Inspects the given {@link ApplicationContext} for {@link MongoPersistentEntityIndexCreator} and those in turn if
 	 * they were registered for the current {@link MappingContext}. If no creator for the current {@link MappingContext}
@@ -1131,8 +1181,22 @@ protected long doCount(String collectionName, Document filter, CountOptions opti
 					.debug(String.format("Executing count: %s in collection: %s", serializeToJsonSafely(filter), collectionName));
 		}
 
-		return execute(collectionName,
-				collection -> collection.countDocuments(CountQuery.of(filter).toQueryDocument(), options));
+		return countExecution.countDocuments(collectionName, filter, options);
+	}
+
+	protected boolean countCanBeEstimated(Document filter, CountOptions options) {
+
+		return
+		// only empty filter for estimatedCount
+		filter.isEmpty() &&
+		// no skip, no limit,...
+				isEmptyOptions(options) &&
+				// transaction active?
+				!MongoDatabaseUtils.isTransactionActive(getMongoDatabaseFactory());
+	}
+
+	private boolean isEmptyOptions(CountOptions options) {
+		return options.getLimit() <= 0 && options.getSkip() <= 0;
 	}
 
 	/*
@@ -3569,5 +3633,15 @@ public MongoDatabase getDb() {
 			// native MongoDB objects that offer methods with ClientSession must not be proxied.
 			return delegate.getDb();
 		}
+
+		@Override
+		protected boolean countCanBeEstimated(Document filter, CountOptions options) {
+			return false;
+		}
+	}
+
+	@FunctionalInterface
+	interface CountExecution {
+		long countDocuments(String collection, Document filter, CountOptions options);
 	}
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java

@@ -2289,6 +2289,34 @@ void createCollectionShouldSetUpTimeSeries() {
 						.granularity(TimeSeriesGranularity.HOURS).toString());
 	}
 
+	@Test // GH-3522
+	void usedCountDocumentsForEmptyQueryByDefault() {
+
+		template.count(new Query(), Human.class);
+
+		verify(collection).countDocuments(any(Document.class), any());
+	}
+
+	@Test // GH-3522
+	void delegatesToEstimatedCountForEmptyQueryIfEnabled() {
+
+		template.useEstimatedCount(true);
+
+		template.count(new Query(), Human.class);
+
+		verify(collection).estimatedDocumentCount(any());
+	}
+
+	@Test // GH-3522
+	void stillUsesCountDocumentsForNonEmptyQueryEvenIfEstimationEnabled() {
+
+		template.useEstimatedCount(true);
+
+		template.count(new BasicQuery("{ 'spring' : 'data-mongodb' }"), Human.class);
+
+		verify(collection).countDocuments(any(Document.class), any());
+	}
+
 	class AutogenerateableId {
 
 		@Id BigInteger id;

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/SessionBoundMongoTemplateUnitTests.java

@@ -220,7 +220,7 @@ public void collectionExistsShouldUseProxiedDatabase() {
 		verify(database).listCollectionNames(eq(clientSession));
 	}
 
-	@Test // DATAMONGO-1880
+	@Test // DATAMONGO-1880, GH-3522
 	public void countShouldUseProxiedCollection() {
 
 		template.count(new Query(), Person.class);

src/main/asciidoc/reference/mongodb.adoc

@@ -2152,6 +2152,11 @@ So in version 2.x `MongoOperations.count()` would use the collection statistics
 As of Spring Data MongoDB 3.x any `count` operation uses regardless the existence of filter criteria the aggregation-based count approach via MongoDBs `countDocuments`.
 If the application is fine with the limitations of working upon collection statistics `MongoOperations.estimatedCount()` offers an alternative.
 
+[TIP]
+====
+By setting `MongoTemplate#useEstimatedCount(...)` to `true` _count_ operations, that use an empty filter query, will be delegated to `estimatedCount`, as long as there is no transaction active and the template is not bound to a <<mongo.sessions,session>>.
+====
+
 [NOTE]
 ====
 MongoDBs native `countDocuments` method and the `$match` aggregation, do not support `$near` and `$nearSphere` but require `$geoWithin` along with `$center` or `$centerSphere` which does not support `$minDistance` (see https://jira.mongodb.org/browse/SERVER-37043).