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.
To still be able to retrieve the exact number of matching documents we also introduced MongoTemplate#exactCount.

Closes: #3522
Original pull request: #3951.

Author: christophstrobl

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

@@ -1179,8 +1179,11 @@ <S, T> T findAndReplace(Query query, S replacement, FindAndReplaceOptions option
 	 *          {@literal null}.
 	 * @param entityClass class that determines the collection to use. Must not be {@literal null}.
 	 * @return the count of matching documents.
+	 * @since 3.4
 	 */
-	long count(Query query, Class<?> entityClass);
+	default long exactCount(Query query, Class<?> entityClass) {
+		return exactCount(query, entityClass, getCollectionName(entityClass));
+	}
 
 	/**
 	 * Returns the number of documents for the given {@link Query} querying the given collection. The given {@link Query}
@@ -1201,6 +1204,71 @@ <S, T> T findAndReplace(Query query, S replacement, FindAndReplaceOptions option
 	 * @param collectionName must not be {@literal null} or empty.
 	 * @return the count of matching documents.
 	 * @see #count(Query, Class, String)
+	 * @since 3.4
+	 */
+	default long exactCount(Query query, String collectionName) {
+		return exactCount(query, null, collectionName);
+	}
+
+	/**
+	 * Returns the number of documents for the given {@link Query} by querying the given collection using the given entity
+	 * class to map the given {@link Query}. <br />
+	 * <strong>NOTE:</strong> Query {@link Query#getSkip() offset} and {@link Query#getLimit() limit} can have direct
+	 * influence on the resulting number of documents found as those values are passed on to the server and potentially
+	 * limit the range and order within which the server performs the count operation. Use an {@literal unpaged} query to
+	 * count all matches.
+	 * <br />
+	 * This method uses an
+	 * {@link com.mongodb.client.MongoCollection#countDocuments(org.bson.conversions.Bson, com.mongodb.client.model.CountOptions)
+	 * aggregation execution} even for empty {@link Query queries} which may have an impact on performance, but guarantees
+	 * shard, session and transaction compliance. In case an inaccurate count satisfies the applications needs use
+	 * {@link #estimatedCount(String)} for empty queries instead.
+	 *
+	 * @param query the {@link Query} class that specifies the criteria used to find documents. Must not be
+	 *          {@literal null}.
+	 * @param entityClass the parametrized type. Can be {@literal null}.
+	 * @param collectionName must not be {@literal null} or empty.
+	 * @return the count of matching documents.
+	 * @since 3.4
+	 */
+	long exactCount(Query query, @Nullable Class<?> entityClass, String collectionName);
+
+	/**
+	 * Returns the number of documents for the given {@link Query} by querying the collection of the given entity class.
+	 * <br />
+	 * <strong>NOTE:</strong> Query {@link Query#getSkip() offset} and {@link Query#getLimit() limit} can have direct
+	 * influence on the resulting number of documents found as those values are passed on to the server and potentially
+	 * limit the range and order within which the server performs the count operation. Use an {@literal unpaged} query to
+	 * count all matches.
+	 * <br />
+	 * This method may choose to use {@link #estimatedCount(Class)} for empty queries instead of running an
+	 * {@link com.mongodb.client.MongoCollection#countDocuments(org.bson.conversions.Bson, com.mongodb.client.model.CountOptions)
+	 * aggregation execution} which may have an impact on performance.
+	 *
+	 * @param query the {@link Query} class that specifies the criteria used to find documents. Must not be
+	 *          {@literal null}.
+	 * @param entityClass class that determines the collection to use. Must not be {@literal null}.
+	 * @return the count of matching documents.
+	 */
+	long count(Query query, Class<?> entityClass);
+
+	/**
+	 * Returns the number of documents for the given {@link Query} querying the given collection. The given {@link Query}
+	 * must solely consist of document field references as we lack type information to map potential property references
+	 * onto document fields. Use {@link #count(Query, Class, String)} to get full type specific support. <br />
+	 * <strong>NOTE:</strong> Query {@link Query#getSkip() offset} and {@link Query#getLimit() limit} can have direct
+	 * influence on the resulting number of documents found as those values are passed on to the server and potentially
+	 * limit the range and order within which the server performs the count operation. Use an {@literal unpaged} query to
+	 * count all matches.
+	 * <br />
+	 * This method may choose to use {@link #estimatedCount(Class)} for empty queries instead of running an
+	 * {@link com.mongodb.client.MongoCollection#countDocuments(org.bson.conversions.Bson, com.mongodb.client.model.CountOptions)
+	 * aggregation execution} which may have an impact on performance.
+	 *
+	 * @param query the {@link Query} class that specifies the criteria used to find documents.
+	 * @param collectionName must not be {@literal null} or empty.
+	 * @return the count of matching documents.
+	 * @see #count(Query, Class, String)
 	 */
 	long count(Query query, String collectionName);
 
@@ -1241,11 +1309,9 @@ default long estimatedCount(Class<?> entityClass) {
 	 * limit the range and order within which the server performs the count operation. Use an {@literal unpaged} query to
 	 * count all matches.
 	 * <br />
-	 * This method uses an
+	 * This method may choose to use {@link #estimatedCount(Class)} for empty queries instead of running an
 	 * {@link com.mongodb.client.MongoCollection#countDocuments(org.bson.conversions.Bson, com.mongodb.client.model.CountOptions)
-	 * aggregation execution} even for empty {@link Query queries} which may have an impact on performance, but guarantees
-	 * shard, session and transaction compliance. In case an inaccurate count satisfies the applications needs use
-	 * {@link #estimatedCount(String)} for empty queries instead.
+	 * aggregation execution} which may have an impact on performance.
 	 *
 	 * @param query the {@link Query} class that specifies the criteria used to find documents. Must not be
 	 *          {@literal null}.

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,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 
 	private SessionSynchronization sessionSynchronization = SessionSynchronization.ON_ACTUAL_TRANSACTION;
 
+	private CountExecution countExecution = this::doExactCount;
+
 	/**
 	 * Constructor used for a basic template configuration.
 	 *
@@ -345,6 +347,47 @@ 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
+	 */
+	private void useEstimatedCount(boolean enabled, BiPredicate<Document, CountOptions> estimationFilter) {
+
+		if (enabled) {
+
+			this.countExecution = (collectionName, filter, options) -> {
+
+				if (!estimationFilter.test(filter, options)) {
+					return doExactCount(collectionName, filter, 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 = this::doExactCount;
+		}
+	}
+
 	/**
 	 * 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}
@@ -1106,6 +1149,17 @@ public long count(Query query, String collectionName) {
 		return count(query, null, collectionName);
 	}
 
+	@Override
+	public long exactCount(Query query, @Nullable Class<?> entityClass, String collectionName) {
+
+		CountContext countContext = queryOperations.countQueryContext(query);
+
+		CountOptions options = countContext.getCountOptions(entityClass);
+		Document mappedQuery = countContext.getMappedQuery(entityClass, mappingContext::getPersistentEntity);
+
+		return doExactCount(collectionName, mappedQuery, options);
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.mongodb.core.MongoOperations#count(org.springframework.data.mongodb.core.query.Query, java.lang.Class, java.lang.String)
@@ -1131,10 +1185,29 @@ protected long doCount(String collectionName, Document filter, CountOptions opti
 					.debug(String.format("Executing count: %s in collection: %s", serializeToJsonSafely(filter), collectionName));
 		}
 
+		return countExecution.countDocuments(collectionName, filter, options);
+	}
+
+	protected long doExactCount(String collectionName, Document filter, CountOptions options) {
 		return execute(collectionName,
 				collection -> collection.countDocuments(CountQuery.of(filter).toQueryDocument(), 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;
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.mongodb.core.MongoOperations#estimatedCount(java.lang.String)
@@ -3571,5 +3644,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/main/java/org/springframework/data/mongodb/core/ReactiveMongoOperations.java

@@ -951,8 +951,11 @@ <S, T> Mono<T> findAndReplace(Query query, S replacement, FindAndReplaceOptions
 	 *          {@literal null}.
 	 * @param entityClass class that determines the collection to use. Must not be {@literal null}.
 	 * @return the count of matching documents.
+	 * @since 3.4
 	 */
-	Mono<Long> count(Query query, Class<?> entityClass);
+	default Mono<Long> exactCount(Query query, Class<?> entityClass) {
+		return exactCount(query, entityClass, getCollectionName(entityClass));
+	}
 
 	/**
 	 * Returns the number of documents for the given {@link Query} querying the given collection. The given {@link Query}
@@ -973,8 +976,11 @@ <S, T> Mono<T> findAndReplace(Query query, S replacement, FindAndReplaceOptions
 	 * @param collectionName must not be {@literal null} or empty.
 	 * @return the count of matching documents.
 	 * @see #count(Query, Class, String)
+	 * @since 3.4
 	 */
-	Mono<Long> count(Query query, String collectionName);
+	default Mono<Long> exactCount(Query query, String collectionName) {
+		return exactCount(query, null, collectionName);
+	}
 
 	/**
 	 * Returns the number of documents for the given {@link Query} by querying the given collection using the given entity
@@ -995,6 +1001,66 @@ <S, T> Mono<T> findAndReplace(Query query, S replacement, FindAndReplaceOptions
 	 * @param entityClass the parametrized type. Can be {@literal null}.
 	 * @param collectionName must not be {@literal null} or empty.
 	 * @return the count of matching documents.
+	 * @since 3.4
+	 */
+	Mono<Long> exactCount(Query query, @Nullable Class<?> entityClass, String collectionName);
+
+	/**
+	 * Returns the number of documents for the given {@link Query} by querying the collection of the given entity class.
+	 * <br />
+	 * <strong>NOTE:</strong> Query {@link Query#getSkip() offset} and {@link Query#getLimit() limit} can have direct
+	 * influence on the resulting number of documents found as those values are passed on to the server and potentially
+	 * limit the range and order within which the server performs the count operation. Use an {@literal unpaged} query to
+	 * count all matches.
+	 * <br />
+	 * This method may choose to use {@link #estimatedCount(Class)} for empty queries instead of running an
+	 * {@link com.mongodb.reactivestreams.client.MongoCollection#countDocuments(org.bson.conversions.Bson, com.mongodb.client.model.CountOptions)
+	 * aggregation execution} which may have an impact on performance.
+	 *
+	 * @param query the {@link Query} class that specifies the criteria used to find documents. Must not be
+	 *          {@literal null}.
+	 * @param entityClass class that determines the collection to use. Must not be {@literal null}.
+	 * @return the count of matching documents.
+	 */
+	Mono<Long> count(Query query, Class<?> entityClass);
+
+	/**
+	 * Returns the number of documents for the given {@link Query} querying the given collection. The given {@link Query}
+	 * must solely consist of document field references as we lack type information to map potential property references
+	 * onto document fields. Use {@link #count(Query, Class, String)} to get full type specific support. <br />
+	 * <strong>NOTE:</strong> Query {@link Query#getSkip() offset} and {@link Query#getLimit() limit} can have direct
+	 * influence on the resulting number of documents found as those values are passed on to the server and potentially
+	 * limit the range and order within which the server performs the count operation. Use an {@literal unpaged} query to
+	 * count all matches.
+	 * <br />
+	 * This method may choose to use {@link #estimatedCount(Class)} for empty queries instead of running an
+	 * {@link com.mongodb.reactivestreams.client.MongoCollection#countDocuments(org.bson.conversions.Bson, com.mongodb.client.model.CountOptions)
+	 * aggregation execution} which may have an impact on performance.
+	 *
+	 * @param query the {@link Query} class that specifies the criteria used to find documents.
+	 * @param collectionName must not be {@literal null} or empty.
+	 * @return the count of matching documents.
+	 * @see #count(Query, Class, String)
+	 */
+	Mono<Long> count(Query query, String collectionName);
+
+	/**
+	 * Returns the number of documents for the given {@link Query} by querying the given collection using the given entity
+	 * class to map the given {@link Query}. <br />
+	 * <strong>NOTE:</strong> Query {@link Query#getSkip() offset} and {@link Query#getLimit() limit} can have direct
+	 * influence on the resulting number of documents found as those values are passed on to the server and potentially
+	 * limit the range and order within which the server performs the count operation. Use an {@literal unpaged} query to
+	 * count all matches.
+	 * <br />
+	 * This method may choose to use {@link #estimatedCount(Class)} for empty queries instead of running an
+	 * {@link com.mongodb.reactivestreams.client.MongoCollection#countDocuments(org.bson.conversions.Bson, com.mongodb.client.model.CountOptions)
+	 * aggregation execution} which may have an impact on performance.
+	 *
+	 * @param query the {@link Query} class that specifies the criteria used to find documents. Must not be
+	 *          {@literal null}.
+	 * @param entityClass the parametrized type. Can be {@literal null}.
+	 * @param collectionName must not be {@literal null} or empty.
+	 * @return the count of matching documents.
 	 */
 	Mono<Long> count(Query query, @Nullable Class<?> entityClass, String collectionName);