CSHARP-2659: Adding support for $sample aggregation pipeline stage. (#1348)

Co-authored-by: Kevin Smith <[email protected]>

Author: JamesKovacs

src/MongoDB.Driver/AggregateFluent.cs

@@ -267,6 +267,11 @@ public override IAggregateFluent<TNewResult> ReplaceWith<TNewResult>(AggregateEx
             return WithPipeline(_pipeline.ReplaceWith(newRoot));
         }
 
+        public override IAggregateFluent<TResult> Sample(long size)
+        {
+            return WithPipeline(_pipeline.Sample(size));
+        }
+
         public override IAggregateFluent<TResult> Search(
             SearchDefinition<TResult> searchDefinition,
             SearchHighlightOptions<TResult> highlight = null,

src/MongoDB.Driver/AggregateFluentBase.cs

@@ -237,6 +237,9 @@ public virtual IAggregateFluent<TNewResult> ReplaceWith<TNewResult>(AggregateExp
             throw new NotImplementedException();
         }
 
+        /// <inheritdoc />
+        public virtual IAggregateFluent<TResult> Sample(long size) => throw new NotImplementedException();
+
         /// <inheritdoc />
         public virtual IAggregateFluent<TResult> Search(
             SearchDefinition<TResult> searchDefinition,

src/MongoDB.Driver/IAggregateFluent.cs

@@ -381,6 +381,13 @@ IAggregateFluent<TNewResult> Lookup<TForeignDocument, TAsElement, TAs, TNewResul
         /// <returns>The fluent aggregate interface.</returns>
         IAggregateFluent<TNewResult> ReplaceWith<TNewResult>(AggregateExpressionDefinition<TResult, TNewResult> newRoot);
 
+        /// <summary>
+        /// Appends a sample stage to the pipeline.
+        /// </summary>
+        /// <param name="size">The sample size.</param>
+        /// <returns>The fluent aggregate interface.</returns>
+        IAggregateFluent<TResult> Sample(long size);
+
         /// <summary>
         /// Appends a $set stage to the pipeline.
         /// </summary>

src/MongoDB.Driver/PipelineDefinitionBuilder.cs

@@ -1097,6 +1097,24 @@ public static PipelineDefinition<TInput, TOutput> ReplaceWith<TInput, TIntermedi
             return pipeline.AppendStage(PipelineStageDefinitionBuilder.ReplaceWith(newRoot, translationOptions));
         }
 
+        /// <summary>
+        /// Appends a $sample stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="size">The sample size.</param>
+        /// <returns>
+        /// A new pipeline with an additional stage.
+        /// </returns>
+        public static PipelineDefinition<TInput, TOutput> Sample<TInput, TOutput>(
+            this PipelineDefinition<TInput, TOutput> pipeline,
+            long size)
+        {
+            Ensure.IsNotNull(pipeline, nameof(pipeline));
+            return pipeline.AppendStage(PipelineStageDefinitionBuilder.Sample<TOutput>(size));
+        }
+
         /// <summary>
         /// Appends a $search stage to the pipeline.
         /// </summary>

src/MongoDB.Driver/PipelineStageDefinitionBuilder.cs

@@ -1757,6 +1757,19 @@ public static PipelineStageDefinition<TInput, TInput> Sort<TInput>(
             return new SortPipelineStageDefinition<TInput>(sort);
         }
 
+        /// <summary>
+        /// Creates a $sample stage.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <param name="size">The size.</param>
+        /// <returns>The stage.</returns>
+        public static PipelineStageDefinition<TInput, TInput> Sample<TInput>(
+            long size)
+        {
+            Ensure.IsGreaterThanOrEqualToZero(size, nameof(size));
+            return new BsonDocumentPipelineStageDefinition<TInput, TInput>(new BsonDocument("$sample", new BsonDocument("size", size)));
+        }
+
         /// <summary>
         /// Creates a $sortByCount stage.
         /// </summary>

tests/MongoDB.Driver.Tests/PipelineDefinitionBuilderTests.cs

@@ -158,6 +158,42 @@ public void Out_with_time_series_options_should_add_expected_stage()
             stages[0].Should().Be("{ $out: { db: 'database', coll: 'collection', timeseries: { timeField: 'time', metaField: 'symbol' } } }");
         }
 
+        [Theory]
+        [InlineData(0)]
+        [InlineData(15)]
+        public void Sample_should_add_expected_stage(long size)
+        {
+            var pipeline = new EmptyPipelineDefinition<BsonDocument>();
+
+            var result = pipeline.Sample(size);
+
+            var stages = RenderStages(result, BsonDocumentSerializer.Instance);
+            stages.Count.Should().Be(1);
+            stages[0].Should().Be("{ $sample: { size: " + size + " } }");
+        }
+
+        [Fact]
+        public void Sample_should_throw_when_pipeline_is_null()
+        {
+            PipelineDefinition<BsonDocument, BsonDocument> pipeline = null;
+
+            var exception = Record.Exception(() => pipeline.Sample(15));
+
+            exception.Should().BeOfType<ArgumentNullException>()
+                .Which.ParamName.Should().Be("pipeline");
+        }
+
+        [Fact]
+        public void Sample_should_throw_when_size_is_negative()
+        {
+            var pipeline = new EmptyPipelineDefinition<BsonDocument>();
+
+            var exception = Record.Exception(() => pipeline.Sample(-1));
+
+            exception.Should().BeOfType<ArgumentOutOfRangeException>()
+                .Which.ParamName.Should().Be("size");
+        }
+
         [Fact]
         public void Search_should_add_expected_stage()
         {