SERVER-25584 Ensure $facet and $bucketAuto run on merging shard.

2 files changed, 23 insertions, 2 deletions

diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index cce0c010f34..0fd8e69d789 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -2003,7 +2003,7 @@ private:
  * The $bucketAuto stage takes a user-specified number of buckets and automatically determines
  * boundaries such that the values are approximately equally distributed between those buckets.
  */
-class DocumentSourceBucketAuto final : public DocumentSource {
+class DocumentSourceBucketAuto final : public DocumentSource, public SplittableDocumentSource {
 public:
     Value serialize(bool explain = false) const final;
     GetDepsReturn getDependencies(DepsTracker* deps) const final;
@@ -2011,6 +2011,16 @@ public:
     void dispose() final;
     const char* getSourceName() const final;
 
+    /**
+     * The $bucketAuto stage must be run on the merging shard.
+     */
+    boost::intrusive_ptr<DocumentSource> getShardSource() final {
+        return nullptr;
+    }
+    boost::intrusive_ptr<DocumentSource> getMergeSource() final {
+        return this;
+    }
+
     static const uint64_t kMaxMemoryUsageBytes = 100 * 1024 * 1024;
 
     static boost::intrusive_ptr<DocumentSourceBucketAuto> create(
diff --git a/src/mongo/db/pipeline/document_source_facet.h b/src/mongo/db/pipeline/document_source_facet.h
index a6ac631c327..267ebe58bb2 100644
--- a/src/mongo/db/pipeline/document_source_facet.h
+++ b/src/mongo/db/pipeline/document_source_facet.h
@@ -57,7 +57,8 @@ class Pipeline;
  * TODO SERVER-24154: Should inherit from SplittableDocumentSource so that it can split in a sharded
  * cluster.
  */
-class DocumentSourceFacet final : public DocumentSourceNeedsMongod {
+class DocumentSourceFacet final : public DocumentSourceNeedsMongod,
+                                  public SplittableDocumentSource {
 public:
     static boost::intrusive_ptr<DocumentSource> createFromBson(
         BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& pExpCtx);
@@ -95,6 +96,16 @@ public:
      */
     void setSource(DocumentSource* source) final;
 
+    /**
+     * The $facet stage must be run on the merging shard.
+     */
+    boost::intrusive_ptr<DocumentSource> getShardSource() final {
+        return nullptr;
+    }
+    boost::intrusive_ptr<DocumentSource> getMergeSource() final {
+        return this;
+    }
+
     // The following are overridden just to forward calls to sub-pipelines.
     void addInvolvedCollections(std::vector<NamespaceString>* collections) const final;
     void doInjectMongodInterface(std::shared_ptr<MongodInterface> mongod) final;