SERVER-533 Add aggregation stage to randomly sample documents

2 files changed, 41 insertions, 0 deletions

diff --git a/jstests/aggregation/bugs/server533.js b/jstests/aggregation/bugs/server533.js
new file mode 100644
index 00000000000..4f3d364c05b
--- /dev/null
+++ b/jstests/aggregation/bugs/server533.js
@@ -0,0 +1,32 @@
+// SERVER-533: Aggregation stage to randomly sample documents.
+
+// For assertErrorCode.
+load('jstests/aggregation/extras/utils.js');
+
+(function() {
+    'use strict';
+
+    var coll = db.agg_sample;
+    coll.drop();
+
+    var nItems = 3;
+    for (var i = 0; i < nItems; i++) {
+        assert.writeOK(coll.insert({_id: i}));
+    }
+
+    [0, 1, nItems, nItems + 1].forEach(function(size) {
+        var results = coll.aggregate([{$sample: {size: size}}]).toArray();
+        assert.eq(results.length, Math.min(size, nItems));
+    });
+
+    // Multiple $sample stages are allowed.
+    var results = coll.aggregate([{$sample: {size: nItems}}, {$sample: {size: 1}}]).toArray();
+    assert.eq(results.length, 1);
+
+    // Invalid options.
+    assertErrorCode(coll, [{$sample: 'string'}], 28745);
+    assertErrorCode(coll, [{$sample: {size: 'string'}}], 28746);
+    assertErrorCode(coll, [{$sample: {size: -1}}], 28747);
+    assertErrorCode(coll, [{$sample: {unknownOpt: true}}], 28748);
+    assertErrorCode(coll, [{$sample: {/* no size */}}], 28749);
+}());
diff --git a/jstests/aggregation/testshard1.js b/jstests/aggregation/testshard1.js
index 71271ca0af5..ab49ad5653a 100644
--- a/jstests/aggregation/testshard1.js
+++ b/jstests/aggregation/testshard1.js
@@ -226,6 +226,15 @@ function testAvgStdDev() {
 }
 testAvgStdDev();
 
+function testSample() {
+    jsTestLog('testing $sample');
+    [0, 1, 10, nItems, nItems + 1].forEach(function(size) {
+        var res = db.ts1.aggregate([{$sample: {size: size}}]).toArray();
+        assert.eq(res.length, Math.min(nItems, size));
+    });
+}
+testSample();
+
 jsTestLog('test $out by copying source collection verbatim to output');
 var outCollection = db.ts1_out;
 var res = aggregateOrdered(db.ts1, [{$out: outCollection.getName()}]);