SERVER-48848 Add support for $sample aggregation stage in grammar and CST

4 files changed, 16 insertions, 6 deletions

diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index 9995e700367..71f7524be9c 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -389,6 +389,7 @@ env.CppUnitTest(
     LIBDEPS=[
         '$BUILD_DIR/mongo/base',
         '$BUILD_DIR/mongo/db/auth/authmocks',
+        '$BUILD_DIR/mongo/db/cst/cst',
         '$BUILD_DIR/mongo/db/exec/document_value/document_value',
         '$BUILD_DIR/mongo/db/exec/document_value/document_value_test_util',
         '$BUILD_DIR/mongo/db/query/collation/collator_interface_mock',
diff --git a/src/mongo/db/pipeline/document_source_sample.cpp b/src/mongo/db/pipeline/document_source_sample.cpp
index 0904778c871..5fd5a99f43e 100644
--- a/src/mongo/db/pipeline/document_source_sample.cpp
+++ b/src/mongo/db/pipeline/document_source_sample.cpp
@@ -92,17 +92,15 @@ const BSONObj randSortSpec = BSON("$rand" << BSON("$meta"
 intrusive_ptr<DocumentSource> DocumentSourceSample::createFromBson(
     BSONElement specElem, const intrusive_ptr<ExpressionContext>& expCtx) {
     uassert(28745, "the $sample stage specification must be an object", specElem.type() == Object);
-    intrusive_ptr<DocumentSourceSample> sample(new DocumentSourceSample(expCtx));
 
     bool sizeSpecified = false;
+    long long size;
     for (auto&& elem : specElem.embeddedObject()) {
         auto fieldName = elem.fieldNameStringData();
 
         if (fieldName == "size") {
             uassert(28746, "size argument to $sample must be a number", elem.isNumber());
-            auto size = elem.safeNumberLong();
-            uassert(28747, "size argument to $sample must not be negative", size >= 0);
-            sample->_size = size;
+            size = elem.safeNumberLong();
             sizeSpecified = true;
         } else {
             uasserted(28748, str::stream() << "unrecognized option to $sample: " << fieldName);
@@ -110,8 +108,16 @@ intrusive_ptr<DocumentSource> DocumentSourceSample::createFromBson(
     }
     uassert(28749, "$sample stage must specify a size", sizeSpecified);
 
-    sample->_sortStage = DocumentSourceSort::create(expCtx, randSortSpec, sample->_size);
+    return DocumentSourceSample::create(expCtx, size);
+}
 
+boost::intrusive_ptr<DocumentSource> DocumentSourceSample::create(
+    const boost::intrusive_ptr<ExpressionContext>& expCtx, long long size) {
+    uassert(28747, "size argument to $sample must not be negative", size >= 0);
+
+    intrusive_ptr<DocumentSourceSample> sample(new DocumentSourceSample(expCtx));
+    sample->_size = size;
+    sample->_sortStage = DocumentSourceSort::create(expCtx, randSortSpec, sample->_size);
     return sample;
 }
 
diff --git a/src/mongo/db/pipeline/document_source_sample.h b/src/mongo/db/pipeline/document_source_sample.h
index 0d10552ca4c..084b4385fc6 100644
--- a/src/mongo/db/pipeline/document_source_sample.h
+++ b/src/mongo/db/pipeline/document_source_sample.h
@@ -67,6 +67,9 @@ public:
     static boost::intrusive_ptr<DocumentSource> createFromBson(
         BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& expCtx);
 
+    static boost::intrusive_ptr<DocumentSource> create(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, long long size);
+
 private:
     explicit DocumentSourceSample(const boost::intrusive_ptr<ExpressionContext>& pExpCtx);
 
diff --git a/src/mongo/db/pipeline/document_source_skip.cpp b/src/mongo/db/pipeline/document_source_skip.cpp
index f38036f0152..33fff89bcef 100644
--- a/src/mongo/db/pipeline/document_source_skip.cpp
+++ b/src/mongo/db/pipeline/document_source_skip.cpp
@@ -100,6 +100,7 @@ Pipeline::SourceContainer::iterator DocumentSourceSkip::doOptimizeAt(
 
 intrusive_ptr<DocumentSourceSkip> DocumentSourceSkip::create(
     const intrusive_ptr<ExpressionContext>& pExpCtx, long long nToSkip) {
+    uassert(15956, "Argument to $skip cannot be negative", nToSkip >= 0);
     intrusive_ptr<DocumentSourceSkip> skip(new DocumentSourceSkip(pExpCtx, nToSkip));
     return skip;
 }
@@ -110,7 +111,6 @@ intrusive_ptr<DocumentSource> DocumentSourceSkip::createFromBson(
             str::stream() << "Argument to $skip must be a number not a " << typeName(elem.type()),
             elem.isNumber());
     auto nToSkip = elem.safeNumberLong();
-    uassert(15956, "Argument to $skip cannot be negative", nToSkip >= 0);
 
     return DocumentSourceSkip::create(pExpCtx, nToSkip);
 }