Revert "SERVER-53397 Desugar $setWindowFields partitionBy using $sort"

This reverts commit dc1c3a791f63dfb909ed520aadab66a4b1ce66e9.

13 files changed, 78 insertions, 307 deletions

diff --git a/src/mongo/db/pipeline/document_source.cpp b/src/mongo/db/pipeline/document_source.cpp
index aa0ca4885ea..d33581db906 100644
--- a/src/mongo/db/pipeline/document_source.cpp
+++ b/src/mongo/db/pipeline/document_source.cpp
@@ -79,19 +79,7 @@ void DocumentSource::registerParser(
             it == parserMap.end());
     parserMap[name] = {parser, requiredMinVersion};
 }
-void DocumentSource::registerParser(
-    string name,
-    SimpleParser simpleParser,
-    boost::optional<ServerGlobalParams::FeatureCompatibility::Version> requiredMinVersion) {
 
-    Parser parser =
-        [simpleParser = std::move(simpleParser)](
-            BSONElement stageSpec,
-            const intrusive_ptr<ExpressionContext>& expCtx) -> list<intrusive_ptr<DocumentSource>> {
-        return {simpleParser(std::move(stageSpec), expCtx)};
-    };
-    return registerParser(std::move(name), std::move(parser), std::move(requiredMinVersion));
-}
 bool DocumentSource::hasQuery() const {
     return false;
 }
diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index e4797743036..b9c89f39bc7 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -70,11 +70,8 @@ class Document;
  * LiteParsedDocumentSource. This is used for checks that need to happen before a full parse,
  * such as checks about which namespaces are referenced by this aggregation.
  *
- * 'fullParser' is either a DocumentSource::SimpleParser or a DocumentSource::Parser.
- * In both cases, it takes a BSONElement and an ExpressionContext and returns fully-executable
- * DocumentSource(s), for optimization and execution. In the common case it's a SimpleParser,
- * which returns a single DocumentSource; in the general case it's a Parser, which returns a whole
- * std::list to support "multi-stage aliases" like $bucket.
+ * 'fullParser' takes a BSONElement and an ExpressionContext and returns a fully-executable
+ * DocumentSource. This will be used for optimization and execution.
  *
  * Stages that do not require any special pre-parse checks can use
  * LiteParsedDocumentSourceDefault::parse as their 'liteParser'.
@@ -85,55 +82,57 @@ class Document;
  * REGISTER_DOCUMENT_SOURCE(foo,
  *                          LiteParsedDocumentSourceDefault::parse,
  *                          DocumentSourceFoo::createFromBson);
- */
-#define REGISTER_DOCUMENT_SOURCE(key, liteParser, fullParser) \
-    REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(key, liteParser, fullParser, boost::none, true)
-
-/**
- * Like REGISTER_DOCUMENT_SOURCE, except the parser will only be enabled when FCV >= minVersion.
- * We store minVersion in the parserMap, so that changing FCV at runtime correctly enables/disables
- * the parser.
- */
-#define REGISTER_DOCUMENT_SOURCE_WITH_MIN_VERSION(key, liteParser, fullParser, minVersion) \
-    REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(key, liteParser, fullParser, minVersion, true)
-
-/**
- * Like REGISTER_DOCUMENT_SOURCE_WITH_MIN_VERSION, except you can also specify a condition,
- * evaluated during startup, that decides whether to register the parser.
- *
- * For example, you could check a feature flag, and register the parser only when it's enabled.
  *
- * Note that the condition is evaluated only once, during a MONGO_INITIALIZER. Don't specify
- * a condition that can change at runtime, such as FCV. (Feature flags are ok, because they
- * cannot be toggled at runtime.)
- *
- * This is the most general REGISTER_DOCUMENT_SOURCE* macro, which all others should delegate to.
+ * If your stage is actually an alias which needs to return more than one stage (such as
+ * $sortByCount), you should use the REGISTER_MULTI_STAGE_ALIAS macro instead.
  */
 #define REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(key, liteParser, fullParser, minVersion, ...) \
     MONGO_INITIALIZER(addToDocSourceParserMap_##key)(InitializerContext*) {                  \
         if (!__VA_ARGS__) {                                                                  \
             return;                                                                          \
         }                                                                                    \
+        auto fullParserWrapper = [](BSONElement stageSpec,                                   \
+                                    const boost::intrusive_ptr<ExpressionContext>& expCtx) { \
+            return std::list<boost::intrusive_ptr<DocumentSource>>{                          \
+                (fullParser)(stageSpec, expCtx)};                                            \
+        };                                                                                   \
         LiteParsedDocumentSource::registerParser("$" #key, liteParser);                      \
-        DocumentSource::registerParser("$" #key, fullParser, minVersion);                    \
+        DocumentSource::registerParser("$" #key, fullParserWrapper, minVersion);             \
     }
 
-/**
- * Like REGISTER_DOCUMENT_SOURCE, except the parser is only enabled when test-commands are enabled.
- */
+#define REGISTER_DOCUMENT_SOURCE(key, liteParser, fullParser) \
+    REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(key, liteParser, fullParser, boost::none, true)
+
 #define REGISTER_TEST_DOCUMENT_SOURCE(key, liteParser, fullParser) \
     REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(                        \
         key, liteParser, fullParser, boost::none, ::mongo::getTestCommandsEnabled())
 
+#define REGISTER_DOCUMENT_SOURCE_WITH_MIN_VERSION(key, liteParser, fullParser, minVersion) \
+    REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(key, liteParser, fullParser, minVersion, true)
+
+/**
+ * Registers a multi-stage alias (such as $sortByCount) to have the single name 'key'. When a stage
+ * with name '$key' is found, 'liteParser' will be used to produce a LiteParsedDocumentSource,
+ * while 'fullParser' will be called to construct a vector of DocumentSources. See the comments on
+ * REGISTER_DOCUMENT_SOURCE for more information.
+ *
+ * As an example, if your stage alias looks like {$foo: <args>} and does *not* require any special
+ * pre-parse checks, you should implement a static parser like DocumentSourceFoo::createFromBson(),
+ * and register it like so:
+ * REGISTER_MULTI_STAGE_ALIAS(foo,
+ *                            LiteParsedDocumentSourceDefault::parse,
+ *                            DocumentSourceFoo::createFromBson);
+ */
+#define REGISTER_MULTI_STAGE_ALIAS(key, liteParser, fullParser)                  \
+    MONGO_INITIALIZER(addAliasToDocSourceParserMap_##key)(InitializerContext*) { \
+        LiteParsedDocumentSource::registerParser("$" #key, (liteParser));        \
+        DocumentSource::registerParser("$" #key, (fullParser), boost::none);     \
+    }
+
 class DocumentSource : public RefCountable {
 public:
-    // In general a parser returns a list of DocumentSources, to accomodate "multi-stage aliases"
-    // like $bucket.
     using Parser = std::function<std::list<boost::intrusive_ptr<DocumentSource>>(
         BSONElement, const boost::intrusive_ptr<ExpressionContext>&)>;
-    // But in the common case a parser returns only one DocumentSource.
-    using SimpleParser = std::function<boost::intrusive_ptr<DocumentSource>(
-        BSONElement, const boost::intrusive_ptr<ExpressionContext>&)>;
 
     using ChangeStreamRequirement = StageConstraints::ChangeStreamRequirement;
     using HostTypeRequirement = StageConstraints::HostTypeRequirement;
@@ -368,17 +367,6 @@ public:
         std::string name,
         Parser parser,
         boost::optional<ServerGlobalParams::FeatureCompatibility::Version> requiredMinVersion);
-    /**
-     * Convenience wrapper for the common case, when DocumentSource::Parser returns a list of one
-     * DocumentSource.
-     *
-     * DO NOT call this method directly. Instead, use the REGISTER_DOCUMENT_SOURCE macro defined in
-     * this file.
-     */
-    static void registerParser(
-        std::string name,
-        SimpleParser simpleParser,
-        boost::optional<ServerGlobalParams::FeatureCompatibility::Version> requiredMinVersion);
 
     /**
      * Returns true if the DocumentSource has a query.
diff --git a/src/mongo/db/pipeline/document_source_add_fields.cpp b/src/mongo/db/pipeline/document_source_add_fields.cpp
index 53d012cd258..2e2b4a5a53d 100644
--- a/src/mongo/db/pipeline/document_source_add_fields.cpp
+++ b/src/mongo/db/pipeline/document_source_add_fields.cpp
@@ -70,19 +70,6 @@ intrusive_ptr<DocumentSource> DocumentSourceAddFields::create(
     return addFields;
 }
 
-intrusive_ptr<DocumentSource> DocumentSourceAddFields::create(
-    const FieldPath& fieldPath,
-    const intrusive_ptr<Expression>& expr,
-    const intrusive_ptr<ExpressionContext>& expCtx) {
-
-    const bool isIndependentOfAnyCollection = false;
-    return make_intrusive<DocumentSourceSingleDocumentTransformation>(
-        expCtx,
-        projection_executor::AddFieldsProjectionExecutor::create(expCtx, fieldPath, expr),
-        kStageName,
-        isIndependentOfAnyCollection);
-}
-
 intrusive_ptr<DocumentSource> DocumentSourceAddFields::createFromBson(
     BSONElement elem, const intrusive_ptr<ExpressionContext>& expCtx) {
     const auto specifiedName = elem.fieldNameStringData();
diff --git a/src/mongo/db/pipeline/document_source_add_fields.h b/src/mongo/db/pipeline/document_source_add_fields.h
index 49cb228c73b..5c99a3790bb 100644
--- a/src/mongo/db/pipeline/document_source_add_fields.h
+++ b/src/mongo/db/pipeline/document_source_add_fields.h
@@ -53,16 +53,6 @@ public:
         StringData stageName = kStageName);
 
     /**
-     * Create a stage that binds an expression to a top-level field.
-     *
-     * 'fieldPath' must be a top-level field name (exactly one element; no dots).
-     */
-    static boost::intrusive_ptr<DocumentSource> create(
-        const FieldPath& fieldPath,
-        const boost::intrusive_ptr<Expression>& expr,
-        const boost::intrusive_ptr<ExpressionContext>& expCtx);
-
-    /**
      * Parses a $addFields stage from the user-supplied BSON.
      */
     static boost::intrusive_ptr<DocumentSource> createFromBson(
diff --git a/src/mongo/db/pipeline/document_source_bucket.cpp b/src/mongo/db/pipeline/document_source_bucket.cpp
index d5761bf3c85..73386201792 100644
--- a/src/mongo/db/pipeline/document_source_bucket.cpp
+++ b/src/mongo/db/pipeline/document_source_bucket.cpp
@@ -40,9 +40,9 @@ using boost::intrusive_ptr;
 using std::list;
 using std::vector;
 
-REGISTER_DOCUMENT_SOURCE(bucket,
-                         LiteParsedDocumentSourceDefault::parse,
-                         DocumentSourceBucket::createFromBson);
+REGISTER_MULTI_STAGE_ALIAS(bucket,
+                           LiteParsedDocumentSourceDefault::parse,
+                           DocumentSourceBucket::createFromBson);
 
 namespace {
 intrusive_ptr<ExpressionConstant> getExpressionConstant(ExpressionContext* const expCtx,
diff --git a/src/mongo/db/pipeline/document_source_change_stream.cpp b/src/mongo/db/pipeline/document_source_change_stream.cpp
index 7b80cb18d02..bad92703daa 100644
--- a/src/mongo/db/pipeline/document_source_change_stream.cpp
+++ b/src/mongo/db/pipeline/document_source_change_stream.cpp
@@ -65,9 +65,9 @@ using std::vector;
 // and re-parse the pipeline. To make this work, the 'transformation' stage will serialize itself
 // with the original specification, and all other stages that are created during the alias expansion
 // will not serialize themselves.
-REGISTER_DOCUMENT_SOURCE(changeStream,
-                         DocumentSourceChangeStream::LiteParsed::parse,
-                         DocumentSourceChangeStream::createFromBson);
+REGISTER_MULTI_STAGE_ALIAS(changeStream,
+                           DocumentSourceChangeStream::LiteParsed::parse,
+                           DocumentSourceChangeStream::createFromBson);
 
 constexpr StringData DocumentSourceChangeStream::kDocumentKeyField;
 constexpr StringData DocumentSourceChangeStream::kFullDocumentBeforeChangeField;
diff --git a/src/mongo/db/pipeline/document_source_count.cpp b/src/mongo/db/pipeline/document_source_count.cpp
index 1107e2e6eeb..80a98a5426b 100644
--- a/src/mongo/db/pipeline/document_source_count.cpp
+++ b/src/mongo/db/pipeline/document_source_count.cpp
@@ -43,9 +43,9 @@ using boost::intrusive_ptr;
 using std::list;
 using std::string;
 
-REGISTER_DOCUMENT_SOURCE(count,
-                         LiteParsedDocumentSourceDefault::parse,
-                         DocumentSourceCount::createFromBson);
+REGISTER_MULTI_STAGE_ALIAS(count,
+                           LiteParsedDocumentSourceDefault::parse,
+                           DocumentSourceCount::createFromBson);
 
 list<intrusive_ptr<DocumentSource>> DocumentSourceCount::createFromBson(
     BSONElement elem, const intrusive_ptr<ExpressionContext>& pExpCtx) {
diff --git a/src/mongo/db/pipeline/document_source_project.cpp b/src/mongo/db/pipeline/document_source_project.cpp
index 6ea0e53e055..e5f048a8e3d 100644
--- a/src/mongo/db/pipeline/document_source_project.cpp
+++ b/src/mongo/db/pipeline/document_source_project.cpp
@@ -33,7 +33,6 @@
 
 #include <boost/optional.hpp>
 #include <boost/smart_ptr/intrusive_ptr.hpp>
-#include <memory>
 
 #include "mongo/db/exec/projection_executor.h"
 #include "mongo/db/exec/projection_executor_builder.h"
@@ -101,26 +100,6 @@ intrusive_ptr<DocumentSource> DocumentSourceProject::create(
     return project;
 }
 
-boost::intrusive_ptr<DocumentSource> DocumentSourceProject::createUnset(
-    const FieldPath& fieldPath, const boost::intrusive_ptr<ExpressionContext>& expCtx) {
-
-    // This helper is only meant for removing top-level fields. Dotted field paths require
-    // thinking about implicit array traversal.
-    tassert(5339701,
-            str::stream() << "Expected a top-level field name, but got " << fieldPath.fullPath(),
-            fieldPath.getPathLength() == 1);
-
-    projection_ast::ProjectionPathASTNode pathNode;
-    pathNode.addChild(fieldPath.fullPath(),
-                      std::make_unique<projection_ast::BooleanConstantASTNode>(false));
-    auto projection = projection_ast::Projection{
-        std::move(pathNode),
-        projection_ast::ProjectType::kExclusion,
-    };
-
-    return create(std::move(projection), expCtx, kAliasNameUnset);
-}
-
 intrusive_ptr<DocumentSource> DocumentSourceProject::createFromBson(
     BSONElement elem, const intrusive_ptr<ExpressionContext>& expCtx) {
     if (elem.fieldNameStringData() == kStageName) {
diff --git a/src/mongo/db/pipeline/document_source_project.h b/src/mongo/db/pipeline/document_source_project.h
index 9401231a56b..13d45282a75 100644
--- a/src/mongo/db/pipeline/document_source_project.h
+++ b/src/mongo/db/pipeline/document_source_project.h
@@ -60,24 +60,16 @@ public:
         BSONObj projectSpec,
         const boost::intrusive_ptr<ExpressionContext>& expCtx,
         StringData specifiedName) try {
-
-        auto projection = projection_ast::parse(
-            expCtx, projectSpec, ProjectionPolicies::aggregateProjectionPolicies());
-        return create(projection, expCtx, specifiedName);
+        return create(projection_ast::parse(
+                          expCtx, projectSpec, ProjectionPolicies::aggregateProjectionPolicies()),
+                      expCtx,
+                      specifiedName);
     } catch (DBException& ex) {
         ex.addContext("Invalid " + specifiedName.toString());
         throw;
     }
 
     /**
-     * Create an '$unset' stage, which removes a single top-level field.
-     *
-     * 'fieldPath' must be a top-level field.
-     */
-    static boost::intrusive_ptr<DocumentSource> createUnset(
-        const FieldPath& fieldPath, const boost::intrusive_ptr<ExpressionContext>& expCtx);
-
-    /**
      * Parses a $project stage from the user-supplied BSON.
      */
     static boost::intrusive_ptr<DocumentSource> createFromBson(
diff --git a/src/mongo/db/pipeline/document_source_set_window_fields.cpp b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
index f38d76b6c14..a98069a6c81 100644
--- a/src/mongo/db/pipeline/document_source_set_window_fields.cpp
+++ b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
@@ -29,153 +29,21 @@
 
 #include "mongo/platform/basic.h"
 
-#include "mongo/db/pipeline/document_source_add_fields.h"
-#include "mongo/db/pipeline/document_source_project.h"
 #include "mongo/db/pipeline/document_source_set_window_fields.h"
 #include "mongo/db/pipeline/document_source_set_window_fields_gen.h"
-#include "mongo/db/pipeline/document_source_sort.h"
 #include "mongo/db/pipeline/lite_parsed_document_source.h"
 #include "mongo/db/query/query_feature_flags_gen.h"
 
-using boost::intrusive_ptr;
-using boost::optional;
-using std::list;
-
 namespace mongo {
 
 REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(
     setWindowFields,
     LiteParsedDocumentSourceDefault::parse,
-    document_source_set_window_fields::createFromBson,
-    boost::none,
-    ::mongo::feature_flags::gFeatureFlagWindowFunctions.isEnabledAndIgnoreFCV());
-
-REGISTER_DOCUMENT_SOURCE_CONDITIONALLY(
-    _internalSetWindowFields,
-    LiteParsedDocumentSourceDefault::parse,
-    DocumentSourceInternalSetWindowFields::createFromBson,
+    DocumentSourceSetWindowFields::createFromBson,
     boost::none,
     ::mongo::feature_flags::gFeatureFlagWindowFunctions.isEnabledAndIgnoreFCV());
 
-list<intrusive_ptr<DocumentSource>> document_source_set_window_fields::createFromBson(
-    BSONElement elem, const intrusive_ptr<ExpressionContext>& expCtx) {
-    uassert(ErrorCodes::FailedToParse,
-            str::stream() << "the " << kStageName
-                          << " stage specification must be an object, found "
-                          << typeName(elem.type()),
-            elem.type() == BSONType::Object);
-
-    auto spec =
-        SetWindowFieldsSpec::parse(IDLParserErrorContext(kStageName), elem.embeddedObject());
-    auto partitionBy = [&]() -> boost::optional<boost::intrusive_ptr<Expression>> {
-        if (auto partitionBy = spec.getPartitionBy())
-            return Expression::parseOperand(
-                expCtx.get(), partitionBy->getElement(), expCtx->variablesParseState);
-        else
-            return boost::none;
-    }();
-    return create(std::move(expCtx),
-                  std::move(partitionBy),
-                  std::move(spec.getSortBy()),
-                  std::move(spec.getOutput()));
-}
-
-list<intrusive_ptr<DocumentSource>> document_source_set_window_fields::create(
-    const intrusive_ptr<ExpressionContext>& expCtx,
-    optional<intrusive_ptr<Expression>> partitionBy,
-    optional<BSONObj> sortBy,
-    BSONObj fields) {
-
-    // Starting with an input like this:
-    //     {$setWindowFields: {partitionBy: {$foo: "$x"}, sortBy: {y: 1}, fields: {...}}}
-
-    // We move the partitionBy expression out into its own $set stage:
-    //     {$set: {__tmp: {$foo: "$x"}}}
-    //     {$setWindowFields: {partitionBy: "$__tmp", sortBy: {y: 1}, fields: {...}}}
-    //     {$unset: '__tmp'}
-
-    // This lets us insert a $sort in between:
-    //     {$set: {__tmp: {$foo: "$x"}}}
-    //     {$sort: {__tmp: 1, y: 1}}
-    //     {$setWindowFields: {partitionBy: "$__tmp", sortBy: {y: 1}, fields: {...}}}
-    //     {$unset: '__tmp'}
-
-    // Which lets us replace $setWindowFields with $_internalSetWindowFields:
-    //     {$set: {__tmp: {$foo: "$x"}}}
-    //     {$sort: {__tmp: 1, y: 1}}
-    //     {$_internalSetWindowFields: {partitionBy: "$__tmp", sortBy: {y: 1}, fields: {...}}}
-    //     {$unset: '__tmp'}
-
-    // If partitionBy is a field path, we can $sort by that field directly and avoid creating a
-    // $set stage. This is important for pushing down the $sort. This is only valid because we
-    // assert (in getNextInput()) that partitionBy is never an array.
-
-    // If there is no partitionBy at all then we just $sort by the sortBy spec.
-
-    // If there is no sortBy and no partitionBy then we can omit the $sort stage completely.
-
-    list<intrusive_ptr<DocumentSource>> result;
-
-    // complexPartitionBy is an expression to evaluate.
-    // simplePartitionBy is a field path, which can be evaluated or sorted.
-    optional<intrusive_ptr<Expression>> complexPartitionBy;
-    optional<FieldPath> simplePartitionBy;
-    optional<intrusive_ptr<Expression>> simplePartitionByExpr;
-    // If there is no partitionBy, both are empty.
-    // If partitionBy is already a field path, we only fill in simplePartitionBy.
-    // If partitionBy is a more complex expression, we will need to generate a $set stage,
-    // which will bind the value of the expression to the name in simplePartitionBy.
-    if (partitionBy) {
-        auto exprFieldPath = dynamic_cast<ExpressionFieldPath*>(partitionBy->get());
-        if (exprFieldPath && exprFieldPath->isRootFieldPath()) {
-            // ExpressionFieldPath has "CURRENT" as an explicit first component,
-            // but for $sort we don't want that.
-            simplePartitionBy = exprFieldPath->getFieldPath().tail();
-            simplePartitionByExpr = partitionBy;
-        } else {
-            // In DocumentSource we don't have a mechanism for generating non-colliding field names,
-            // so we have to choose the tmp name carefully to make a collision unlikely in practice.
-            auto tmp = "__internal_setWindowFields_partition_key";
-            simplePartitionBy = FieldPath{tmp};
-            simplePartitionByExpr = ExpressionFieldPath::createPathFromString(
-                expCtx.get(), tmp, expCtx->variablesParseState);
-            complexPartitionBy = partitionBy;
-        }
-    }
-
-    // $set
-    if (complexPartitionBy) {
-        result.push_back(
-            DocumentSourceAddFields::create(*simplePartitionBy, *complexPartitionBy, expCtx));
-    }
-
-    // $sort
-    if (simplePartitionBy || sortBy) {
-        BSONObjBuilder sortSpec;
-        if (simplePartitionBy) {
-            sortSpec << simplePartitionBy->fullPath() << 1;
-        }
-        if (sortBy) {
-            for (auto elem : *sortBy) {
-                sortSpec << elem;
-            }
-        }
-        result.push_back(DocumentSourceSort::create(expCtx, sortSpec.obj()));
-    }
-
-    // $_internalSetWindowFields
-    result.push_back(make_intrusive<DocumentSourceInternalSetWindowFields>(
-        expCtx, simplePartitionByExpr, sortBy, fields));
-
-    // $unset
-    if (complexPartitionBy) {
-        result.push_back(DocumentSourceProject::createUnset(*simplePartitionBy, expCtx));
-    }
-
-    return result;
-}
-
-Value DocumentSourceInternalSetWindowFields::serialize(
+Value DocumentSourceSetWindowFields::serialize(
     boost::optional<ExplainOptions::Verbosity> explain) const {
     MutableDocument spec;
     spec[SetWindowFieldsSpec::kPartitionByFieldName] =
@@ -185,7 +53,7 @@ Value DocumentSourceInternalSetWindowFields::serialize(
     return Value(DOC(kStageName << spec.freeze()));
 }
 
-boost::intrusive_ptr<DocumentSource> DocumentSourceInternalSetWindowFields::createFromBson(
+boost::intrusive_ptr<DocumentSource> DocumentSourceSetWindowFields::createFromBson(
     BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& expCtx) {
     uassert(ErrorCodes::FailedToParse,
             str::stream() << "the " << kStageName
@@ -202,13 +70,12 @@ boost::intrusive_ptr<DocumentSource> DocumentSourceInternalSetWindowFields::crea
         else
             return boost::none;
     }();
-    return make_intrusive<DocumentSourceInternalSetWindowFields>(
+    return make_intrusive<DocumentSourceSetWindowFields>(
         expCtx, partitionBy, spec.getSortBy(), spec.getOutput());
 }
 
-DocumentSource::GetNextResult DocumentSourceInternalSetWindowFields::doGetNext() {
-    // This is a placeholder: it returns every input doc unchanged.
-    return pSource->getNext();
+DocumentSource::GetNextResult DocumentSourceSetWindowFields::doGetNext() {
+    return GetNextResult::makeEOF();
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_set_window_fields.h b/src/mongo/db/pipeline/document_source_set_window_fields.h
index 15902de2211..09210c3df0d 100644
--- a/src/mongo/db/pipeline/document_source_set_window_fields.h
+++ b/src/mongo/db/pipeline/document_source_set_window_fields.h
@@ -34,26 +34,9 @@
 
 namespace mongo {
 
-/**
- * $setWindowFields is an alias: it desugars to some combination of projection, sorting,
- * and $_internalSetWindowFields.
- */
-namespace document_source_set_window_fields {
-constexpr StringData kStageName = "$setWindowFields"_sd;
-
-std::list<boost::intrusive_ptr<DocumentSource>> createFromBson(
-    BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& pExpCtx);
-
-std::list<boost::intrusive_ptr<DocumentSource>> create(
-    const boost::intrusive_ptr<ExpressionContext>& expCtx,
-    boost::optional<boost::intrusive_ptr<Expression>> partitionBy,
-    boost::optional<BSONObj> sortBy,
-    BSONObj fields);
-}  // namespace document_source_set_window_fields
-
-class DocumentSourceInternalSetWindowFields final : public DocumentSource {
+class DocumentSourceSetWindowFields final : public DocumentSource {
 public:
-    static constexpr StringData kStageName = "$_internalSetWindowFields"_sd;
+    static constexpr StringData kStageName = "$setWindowFields"_sd;
 
     /**
      * Parses 'elem' into a $setWindowFields stage, or throws a AssertionException if 'elem' was an
@@ -63,11 +46,10 @@ public:
         BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& pExpCtx);
 
 
-    DocumentSourceInternalSetWindowFields(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx,
-        boost::optional<boost::intrusive_ptr<Expression>> partitionBy,
-        boost::optional<BSONObj> sortBy,
-        BSONObj fields)
+    DocumentSourceSetWindowFields(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                                  boost::optional<boost::intrusive_ptr<Expression>> partitionBy,
+                                  boost::optional<BSONObj> sortBy,
+                                  BSONObj fields)
         : DocumentSource(kStageName, expCtx),
           _partitionBy(partitionBy),
           _sortBy(std::move(sortBy)),
@@ -98,8 +80,6 @@ public:
     DocumentSource::GetNextResult doGetNext();
 
 private:
-    DocumentSource::GetNextResult getNextInput();
-
     boost::optional<boost::intrusive_ptr<Expression>> _partitionBy;
     boost::optional<BSONObj> _sortBy;
     BSONObj _fields;
diff --git a/src/mongo/db/pipeline/document_source_set_window_fields_test.cpp b/src/mongo/db/pipeline/document_source_set_window_fields_test.cpp
index 414e93a53c3..60a2035cb21 100644
--- a/src/mongo/db/pipeline/document_source_set_window_fields_test.cpp
+++ b/src/mongo/db/pipeline/document_source_set_window_fields_test.cpp
@@ -46,47 +46,47 @@ namespace {
 using DocumentSourceSetWindowFieldsTest = AggregationContextFixture;
 
 TEST_F(DocumentSourceSetWindowFieldsTest, FailsToParseInvalidArgumentTypes) {
-    auto spec = BSON("$_internalSetWindowFields"
+    auto spec = BSON("$setWindowFields"
                      << "invalid");
     ASSERT_THROWS_CODE(
-        DocumentSourceInternalSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
+        DocumentSourceSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
         AssertionException,
         ErrorCodes::FailedToParse);
 
-    spec = BSON("$_internalSetWindowFields" << BSON("sortBy"
-                                                    << "invalid sort spec"));
+    spec = BSON("$setWindowFields" << BSON("sortBy"
+                                           << "invalid sort spec"));
     ASSERT_THROWS_CODE(
-        DocumentSourceInternalSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
+        DocumentSourceSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
         AssertionException,
         ErrorCodes::TypeMismatch);
 
-    spec = BSON("$_internalSetWindowFields" << BSON("output"
-                                                    << "invalid"));
+    spec = BSON("$setWindowFields" << BSON("output"
+                                           << "invalid"));
     ASSERT_THROWS_CODE(
-        DocumentSourceInternalSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
+        DocumentSourceSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
         AssertionException,
         ErrorCodes::TypeMismatch);
 
-    spec = BSON("$_internalSetWindowFields"
+    spec = BSON("$setWindowFields"
                 << BSON("partitionBy" << BSON("$notAnExpression" << 1) << "output" << BSONObj()));
     ASSERT_THROWS_CODE(
-        DocumentSourceInternalSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
+        DocumentSourceSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
         AssertionException,
         ErrorCodes::InvalidPipelineOperator);
 
-    spec = BSON("$_internalSetWindowFields" << BSON("unknown_parameter" << 1));
+    spec = BSON("$setWindowFields" << BSON("unknown_parameter" << 1));
     ASSERT_THROWS_CODE(
-        DocumentSourceInternalSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
+        DocumentSourceSetWindowFields::createFromBson(spec.firstElement(), getExpCtx()),
         AssertionException,
         40415);
 }
 
 TEST_F(DocumentSourceSetWindowFieldsTest, SuccessfullyParsesAndReserializes) {
     auto spec = fromjson(R"(
-        {$_internalSetWindowFields: {partitionBy: '$state', sortBy: {city: 1}, output: {mySum: {$sum: 
+        {$setWindowFields: {partitionBy: '$state', sortBy: {city: 1}, output: {mySum: {$sum: 
         {input: '$pop', documents: [-10, 0]}}}}})");
     auto parsedStage =
-        DocumentSourceInternalSetWindowFields::createFromBson(spec.firstElement(), getExpCtx());
+        DocumentSourceSetWindowFields::createFromBson(spec.firstElement(), getExpCtx());
     std::vector<Value> serializedArray;
     parsedStage->serializeToArray(serializedArray);
     ASSERT_BSONOBJ_EQ(serializedArray[0].getDocument().toBson(), spec);
@@ -94,7 +94,7 @@ TEST_F(DocumentSourceSetWindowFieldsTest, SuccessfullyParsesAndReserializes) {
 
 TEST_F(DocumentSourceSetWindowFieldsTest, FailsToParseIfFeatureFlagDisabled) {
     auto spec = fromjson(R"(
-        {$_internalSetWindowFields: {partitionBy: '$state', sortBy: {city: 1}, output: {mySum: {$sum: 
+        {$setWindowFields: {partitionBy: '$state', sortBy: {city: 1}, output: {mySum: {$sum: 
         {input: '$pop', documents: [-10, 0]}}}}})");
     // By default, the unit test will have the feature flag disabled.
     ASSERT_THROWS_CODE(
diff --git a/src/mongo/db/pipeline/document_source_sort_by_count.cpp b/src/mongo/db/pipeline/document_source_sort_by_count.cpp
index 5efce02e3a0..49ea2dbc02a 100644
--- a/src/mongo/db/pipeline/document_source_sort_by_count.cpp
+++ b/src/mongo/db/pipeline/document_source_sort_by_count.cpp
@@ -42,9 +42,9 @@ namespace mongo {
 using boost::intrusive_ptr;
 using std::list;
 
-REGISTER_DOCUMENT_SOURCE(sortByCount,
-                         LiteParsedDocumentSourceDefault::parse,
-                         DocumentSourceSortByCount::createFromBson);
+REGISTER_MULTI_STAGE_ALIAS(sortByCount,
+                           LiteParsedDocumentSourceDefault::parse,
+                           DocumentSourceSortByCount::createFromBson);
 
 list<intrusive_ptr<DocumentSource>> DocumentSourceSortByCount::createFromBson(
     BSONElement elem, const intrusive_ptr<ExpressionContext>& pExpCtx) {