spring-projectsGH-3726: Add support for $sampleRate match expression.

Closes spring-projects#3726

Author: JamesMcNee

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.java

@@ -50,6 +50,7 @@
  * @author Nikolay Bogdanov
  * @author Gustavo de Geus
  * @author Jérôme Guyon
+ * @author James McNee
  * @since 1.3
  */
 public class Aggregation {
@@ -499,6 +500,17 @@ public static MatchOperation match(CriteriaDefinition criteria) {
 		return new MatchOperation(criteria);
 	}
 
+	/**
+	 * Creates a new {@link MatchOperation} using the given {@link MatchExpression}.
+	 *
+	 * @param matchExpression must not be {@literal null}.
+	 * @return new instance of {@link MatchOperation}.
+	 * @since 3.2.4
+	 */
+	public static MatchOperation match(MatchExpression matchExpression) {
+		return new MatchOperation(matchExpression);
+	}
+
 	/**
 	 * Creates a new {@link GeoNearOperation} instance from the given {@link NearQuery} and the {@code distanceField}. The
 	 * {@code distanceField} defines output field that contains the calculated distance.

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/MatchExpression.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2015-2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.aggregation;
+
+import org.bson.Document;
+import org.springframework.data.mongodb.MongoExpression;
+
+/**
+ * A {@link MatchExpression} can be used within the {@code match} aggregation pipeline stage.
+ *
+ * @author James McNee
+ */
+public interface MatchExpression extends MongoExpression {
+
+	/**
+	 * Obtain the as is (unmapped) representation of the {@link MatchExpression}. Use {@link #toDocument(AggregationOperationContext)}
+	 * with a matching {@link AggregationOperationContext context} to engage domain type mapping including field name resolution.
+	 *
+	 * @see org.springframework.data.mongodb.MongoExpression#toDocument()
+	 */
+	@Override
+	default Document toDocument() {
+		return toDocument(Aggregation.DEFAULT_CONTEXT);
+	}
+
+	/**
+	 * Turns the {@link MatchExpression} into a {@link Document} within the given {@link AggregationOperationContext}.
+	 *
+	 * @param context must not be {@literal null}.
+	 * @return the MongoDB native ({@link Document}) form of the expression.
+	 */
+	Document toDocument(AggregationOperationContext context);
+}

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/MatchOperation.java

@@ -29,13 +29,14 @@
  * @author Sebastian Herold
  * @author Thomas Darimont
  * @author Oliver Gierke
+ * @author James McNee
  * @since 1.3
  * @see <a href="https://docs.mongodb.com/manual/reference/operator/aggregation/match/">MongoDB Aggregation Framework:
  *      $match</a>
  */
 public class MatchOperation implements AggregationOperation {
 
-	private final CriteriaDefinition criteriaDefinition;
+	private final Document document;
 
 	/**
 	 * Creates a new {@link MatchOperation} for the given {@link CriteriaDefinition}.
@@ -45,7 +46,19 @@ public class MatchOperation implements AggregationOperation {
 	public MatchOperation(CriteriaDefinition criteriaDefinition) {
 
 		Assert.notNull(criteriaDefinition, "Criteria must not be null!");
-		this.criteriaDefinition = criteriaDefinition;
+		this.document = criteriaDefinition.getCriteriaObject();
+	}
+
+	/**
+	 * Creates a new {@link MatchOperation} for the given {@link MatchExpression}.
+	 *
+	 * @param matchExpression must not be {@literal null}.
+	 */
+	public MatchOperation(MatchExpression matchExpression) {
+
+		Assert.notNull(matchExpression, "Match expression must not be null!");
+		this.document = matchExpression.toDocument();
+
 	}
 
 	/*
@@ -54,7 +67,7 @@ public MatchOperation(CriteriaDefinition criteriaDefinition) {
 	 */
 	@Override
 	public Document toDocument(AggregationOperationContext context) {
-		return new Document(getOperator(), context.getMappedObject(criteriaDefinition.getCriteriaObject()));
+		return new Document(getOperator(), context.getMappedObject(document));
 	}
 
 	/*

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/SamplingOperators.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2016-2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.aggregation;
+
+import org.bson.Document;
+import org.springframework.util.Assert;
+
+/**
+ * Gateway to {@literal sampling expressions} that match documents using a provided criteria.
+ *
+ * @author James McNee
+ * @since 3.2.4
+ */
+public class SamplingOperators {
+
+	/**
+	 * Encapsulates the aggregation framework {@code $sampleRate} operator.
+	 *
+	 * @author James McNee
+	 * @see <a href="https://docs.mongodb.com/manual/reference/operator/aggregation/sampleRate/">
+	 * 		https://docs.mongodb.com/manual/reference/operator/aggregation/sampleRate/</a>
+	 */
+	public static class SampleRate implements MatchExpression {
+
+		private final double sampleRate;
+
+		private SampleRate(double sampleRate) {
+			this.sampleRate = sampleRate;
+		}
+
+		/**
+		 * Creates new {@link SamplingOperators.SampleRate}.
+		 *
+		 * @param sampleRate sample rate to determine number of documents to be randomly selected from the input.
+		 * @return new instance of {@link SamplingOperators.SampleRate}.
+		 */
+		public static SampleRate sampleRate(double sampleRate) {
+			Assert.isTrue(sampleRate >= 0, "Sample rate must be greater than zero!");
+			Assert.isTrue(sampleRate <= 1, "Sample rate must not be greater than one!");
+
+			return new SampleRate(sampleRate);
+		}
+
+		@Override
+		public Document toDocument(AggregationOperationContext context) {
+			return new Document("$sampleRate", sampleRate);
+		}
+	}
+}

Diff for: spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.java

@@ -18,6 +18,7 @@
 import static org.springframework.data.domain.Sort.Direction.*;
 import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
 import static org.springframework.data.mongodb.core.aggregation.Fields.*;
+import static org.springframework.data.mongodb.core.aggregation.SamplingOperators.SampleRate.*;
 import static org.springframework.data.mongodb.core.query.Criteria.*;
 import static org.springframework.data.mongodb.test.util.Assertions.*;
 
@@ -89,6 +90,7 @@
  * @author Maninder Singh
  * @author Sergey Shcherbakov
  * @author Minsu Kim
+ * @author James McNee
  */
 @ExtendWith(MongoTemplateExtension.class)
 public class AggregationTests {
@@ -1944,6 +1946,22 @@ void mapsEnumsInMatchClauseUsingInCriteriaCorrectly() {
 		assertThat(results.getMappedResults()).hasSize(1);
 	}
 
+	@Test // GH-3726
+	void shouldApplySampleRateCorrectly() {
+
+		createUserWithLikesDocuments();
+
+		TypedAggregation<UserWithLikes> agg = newAggregation(UserWithLikes.class,
+				unwind("likes"),
+				match(sampleRate(0.66))
+		);
+
+		assertThat(agg.toString()).isNotNull();
+
+		AggregationResults<LikeStats> result = mongoTemplate.aggregate(agg, LikeStats.class);
+		assertThat(result.getMappedResults().size()).isGreaterThan(0);
+	}
+
 	private void createUsersWithReferencedPersons() {
 
 		mongoTemplate.dropCollection(User.class);

Diff for: spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/SamplingOperatorsUnitTest.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2018-2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.aggregation;
+
+import org.bson.Document;
+import org.junit.jupiter.api.Test;
+import org.springframework.data.mongodb.core.aggregation.SamplingOperators.SampleRate;
+
+import static org.assertj.core.api.Assertions.*;
+
+/**
+ * Unit tests for {@link SamplingOperators}.
+ *
+ * @author James McNee
+ */
+public class SamplingOperatorsUnitTest {
+
+	@Test // GH-3726
+	public void shouldRejectNegativeSampleRate() {
+		assertThatIllegalArgumentException().isThrownBy(() -> SampleRate.sampleRate(-1.0));
+	}
+
+	@Test // GH-3726
+	public void shouldRejectSampleRateGreaterThanOne() {
+		assertThatIllegalArgumentException().isThrownBy(() -> SampleRate.sampleRate(1.1));
+	}
+
+	@Test // GH-3726
+	public void shouldCreateSampleRateExpression() {
+		assertThat(SampleRate.sampleRate(0.34).toDocument()).isEqualTo(Document.parse("{ $sampleRate: 0.34 }"));
+	}
+
+}