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

1 files changed, 12 insertions, 1 deletions

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;