SERVER-42416 Move returnKey execution code to separate PlanStage

1 files changed, 8 insertions, 0 deletions

diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index 48cc9e39027..c8a92a75963 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -52,6 +52,7 @@
 #include "mongo/db/exec/merge_sort.h"
 #include "mongo/db/exec/or.h"
 #include "mongo/db/exec/projection.h"
+#include "mongo/db/exec/return_key.h"
 #include "mongo/db/exec/shard_filter.h"
 #include "mongo/db/exec/skip.h"
 #include "mongo/db/exec/sort.h"
@@ -133,6 +134,13 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx,
             return std::make_unique<SortKeyGeneratorStage>(
                 cq.getExpCtx(), std::move(childStage), ws, keyGenNode->sortSpec);
         }
+        case STAGE_RETURN_KEY: {
+            auto returnKeyNode = static_cast<const ReturnKeyNode*>(root);
+            auto childStage =
+                buildStages(opCtx, collection, cq, qsol, returnKeyNode->children[0], ws);
+            return std::make_unique<ReturnKeyStage>(
+                opCtx, std::move(returnKeyNode->sortKeyMetaFields), ws, std::move(childStage));
+        }
         case STAGE_PROJECTION_DEFAULT: {
             auto pn = static_cast<const ProjectionNodeDefault*>(root);
             auto childStage = buildStages(opCtx, collection, cq, qsol, pn->children[0], ws);