GH-3726: Add support for $sampleRate aggregation operation

Closes #3726

Author: JamesMcNee

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 {
@@ -128,7 +129,7 @@ public static AggregationUpdate newUpdate(AggregationOperation... operations) {
 		return AggregationUpdate.from(Arrays.asList(operations));
 	}
 
-	/**
+    /**
 	 * Returns a copy of this {@link Aggregation} with the given {@link AggregationOptions} set. Note that options are
 	 * supported in MongoDB version 2.6+.
 	 *
@@ -478,6 +479,17 @@ public static SampleOperation sample(long sampleSize) {
 		return new SampleOperation(sampleSize);
 	}
 
+	/**
+	 * Creates a new {@link SampleRateOperation} to select the documents from its input randomly using the provided sample rate.
+	 *
+	 * @param sampleRate must not be less than zero or more than one.
+	 * @return new instance of {@link SampleRateOperation}.
+	 * @since 3.2.4
+	 */
+	public static SampleRateOperation sampleRate(double sampleRate) {
+		return new SampleRateOperation(sampleRate);
+	}
+
 	/**
 	 * Creates a new {@link MatchOperation} using the given {@link Criteria}.
 	 *

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

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2017-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;
+
+/**
+ * Encapsulates the {@code $sampleRate}-operation.
+ * <p>
+ * We recommend using the static factory method {@link Aggregation#sampleRate(double)} instead of creating instances of this
+ * class directly.
+ *
+ * @author James McNee
+ * @see <a href="https://docs.mongodb.com/manual/reference/operator/aggregation/sampleRate/">MongoDB Aggregation Framework:
+ * $sampleRate</a>
+ * @since 3.2.4
+ */
+public class SampleRateOperation implements AggregationOperation {
+
+	private final double sampleRate;
+
+	/**
+	 * @param sampleRate sample rate to determine number of documents to be randomly selected from the input.
+	 */
+	public SampleRateOperation(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!");
+		this.sampleRate = sampleRate;
+	}
+
+	/*
+	  (non-Javadoc)
+	 * @see org.springframework.data.mongodb.core.aggregation.BucketOperationSupport#toDocument(org.springframework.data.mongodb.core.aggregation.AggregationOperationContext)
+	 */
+	@Override
+	public Document toDocument(AggregationOperationContext context) {
+		return new Document(getOperator(), this.sampleRate);
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.core.aggregation.AggregationOperation#getOperator()
+	 */
+	@Override
+	public String getOperator() {
+		return "$sampleRate";
+	}
+}

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2017-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 static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+
+/**
+ * Unit tests for {@link SampleRateOperation}.
+ *
+ * @author James McNee
+ */
+public class SampleRateOperationUnitTests {
+
+	private static final String OPERATION = "$sampleRate";
+
+	@Test // GH-3726
+	public void shouldRejectSampleNegativeSampleRate() {
+		assertThatIllegalArgumentException().isThrownBy(() -> new SampleRateOperation(-1.0));
+	}
+
+	@Test // GH-3726
+	public void shouldRejectSampleRateGreaterThan1() {
+		assertThatIllegalArgumentException().isThrownBy(() -> new SampleRateOperation(1.1));
+	}
+
+	@Test // GH-3726
+	public void rendersSampleRateOperation() {
+		double sampleRate = 0.34;
+
+		SampleRateOperation sampleOperation = Aggregation.sampleRate(sampleRate);
+
+		Document sampleOperationDocument = sampleOperation.toDocument(Aggregation.DEFAULT_CONTEXT);
+
+		assertThat(sampleOperationDocument.get(OPERATION)).isNotNull();
+		assertThat(sampleOperationDocument.get(OPERATION)).isEqualTo(sampleRate);
+	}
+}