summaryrefslogtreecommitdiff
path: root/src/mongo/db/query
diff options
context:
space:
mode:
Diffstat (limited to 'src/mongo/db/query')
-rw-r--r--src/mongo/db/query/get_executor.cpp38
-rw-r--r--src/mongo/db/query/parsed_projection.cpp11
-rw-r--r--src/mongo/db/query/parsed_projection.h6
-rw-r--r--src/mongo/db/query/parsed_projection_test.cpp6
-rw-r--r--src/mongo/db/query/planner_analysis.cpp68
-rw-r--r--src/mongo/db/query/query_planner_common.cpp28
-rw-r--r--src/mongo/db/query/query_planner_common.h8
-rw-r--r--src/mongo/db/query/query_request.cpp17
-rw-r--r--src/mongo/db/query/query_request.h6
-rw-r--r--src/mongo/db/query/query_solution.cpp23
-rw-r--r--src/mongo/db/query/query_solution.h29
-rw-r--r--src/mongo/db/query/stage_builder.cpp8
-rw-r--r--src/mongo/db/query/stage_types.h1
13 files changed, 161 insertions, 88 deletions
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 4b039fa7aaf..4d809846186 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -49,6 +49,7 @@
#include "mongo/db/exec/multi_plan.h"
#include "mongo/db/exec/projection.h"
#include "mongo/db/exec/record_store_fast_count.h"
+#include "mongo/db/exec/return_key.h"
#include "mongo/db/exec/shard_filter.h"
#include "mongo/db/exec/sort_key_generator.h"
#include "mongo/db/exec/subplan.h"
@@ -388,23 +389,32 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx,
std::move(root));
}
- // There might be a projection. The idhack stage will always fetch the full
- // document, so we don't support covered projections. However, we might use the
- // simple inclusion fast path.
- if (nullptr != canonicalQuery->getProj()) {
-
- // Add a SortKeyGeneratorStage if there is a $meta sortKey projection.
- if (canonicalQuery->getProj()->wantSortKey()) {
- root = std::make_unique<SortKeyGeneratorStage>(
- canonicalQuery->getExpCtx(),
- std::move(root),
- ws,
- canonicalQuery->getQueryRequest().getSort());
- }
+ // Add a SortKeyGeneratorStage if there is a $meta sortKey projection.
+ if (canonicalQuery->getProj() && canonicalQuery->getProj()->wantSortKey()) {
+ root = std::make_unique<SortKeyGeneratorStage>(
+ canonicalQuery->getExpCtx(),
+ std::move(root),
+ ws,
+ canonicalQuery->getQueryRequest().getSort());
+ }
+ if (canonicalQuery->getQueryRequest().returnKey()) {
+ // If returnKey was requested, add ReturnKeyStage to return only the index keys in the
+ // resulting documents. If a projection was also specified, it will be ignored, with
+ // the exception the $meta sortKey projection, which can be used along with the
+ // returnKey.
+ root = std::make_unique<ReturnKeyStage>(
+ opCtx,
+ QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(
+ canonicalQuery->getQueryRequest().getProj()),
+ ws,
+ std::move(root));
+ } else if (canonicalQuery->getProj()) {
+ // There might be a projection. The idhack stage will always fetch the full
+ // document, so we don't support covered projections. However, we might use the
+ // simple inclusion fast path.
// Stuff the right data into the params depending on what proj impl we use.
if (canonicalQuery->getProj()->requiresDocument() ||
- canonicalQuery->getProj()->wantIndexKey() ||
canonicalQuery->getProj()->wantSortKey() ||
canonicalQuery->getProj()->hasDottedFieldPath()) {
root = std::make_unique<ProjectionStageDefault>(
diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp
index 359ad5c23d8..165a27e66bb 100644
--- a/src/mongo/db/query/parsed_projection.cpp
+++ b/src/mongo/db/query/parsed_projection.cpp
@@ -56,8 +56,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
IncludeExclude includeExclude = IncludeExclude::kUninitialized;
bool requiresDocument = false;
- bool hasIndexKeyProjection = false;
-
bool wantTextScore = false;
bool wantGeoNearPoint = false;
bool wantGeoNearDistance = false;
@@ -161,7 +159,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
if (e2.valuestr() != QueryRequest::metaTextScore &&
e2.valuestr() != QueryRequest::metaRecordId &&
- e2.valuestr() != QueryRequest::metaIndexKey &&
e2.valuestr() != QueryRequest::metaGeoNearDistance &&
e2.valuestr() != QueryRequest::metaGeoNearPoint &&
e2.valuestr() != QueryRequest::metaSortKey) {
@@ -171,8 +168,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
// This clobbers everything else.
if (e2.valuestr() == QueryRequest::metaTextScore) {
wantTextScore = true;
- } else if (e2.valuestr() == QueryRequest::metaIndexKey) {
- hasIndexKeyProjection = true;
} else if (e2.valuestr() == QueryRequest::metaGeoNearDistance) {
wantGeoNearDistance = true;
} else if (e2.valuestr() == QueryRequest::metaGeoNearPoint) {
@@ -268,7 +263,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
// Save the raw spec. It should be owned by the QueryRequest.
verify(spec.isOwned());
pp->_source = spec;
- pp->_returnKey = hasIndexKeyProjection;
pp->_requiresDocument = requiresDocument;
// Add meta-projections.
@@ -308,11 +302,6 @@ Status ParsedProjection::make(OperationContext* opCtx,
}
}
- // returnKey clobbers everything except for sortKey meta-projection.
- if (hasIndexKeyProjection && !wantSortKey) {
- pp->_requiresDocument = false;
- }
-
*out = pp.release();
return Status::OK();
}
diff --git a/src/mongo/db/query/parsed_projection.h b/src/mongo/db/query/parsed_projection.h
index 0a5ee6de3c2..1d9063acf05 100644
--- a/src/mongo/db/query/parsed_projection.h
+++ b/src/mongo/db/query/parsed_projection.h
@@ -100,10 +100,6 @@ public:
return _wantGeoNearPoint;
}
- bool wantIndexKey() const {
- return _returnKey;
- }
-
bool wantSortKey() const {
return _wantSortKey;
}
@@ -193,8 +189,6 @@ private:
bool _wantGeoNearPoint = false;
- bool _returnKey = false;
-
// Whether this projection includes a sortKey meta-projection.
bool _wantSortKey = false;
diff --git a/src/mongo/db/query/parsed_projection_test.cpp b/src/mongo/db/query/parsed_projection_test.cpp
index 990b665d6ed..8994723a62d 100644
--- a/src/mongo/db/query/parsed_projection_test.cpp
+++ b/src/mongo/db/query/parsed_projection_test.cpp
@@ -234,7 +234,6 @@ TEST(ParsedProjectionTest, ParsedProjectionDefaults) {
ASSERT_FALSE(parsedProjection->requiresMatchDetails());
ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
- ASSERT_FALSE(parsedProjection->wantIndexKey());
}
TEST(ParsedProjectionTest, SortKeyMetaProjection) {
@@ -247,7 +246,6 @@ TEST(ParsedProjectionTest, SortKeyMetaProjection) {
ASSERT_FALSE(parsedProjection->requiresMatchDetails());
ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
- ASSERT_FALSE(parsedProjection->wantIndexKey());
}
TEST(ParsedProjectionTest, SortKeyMetaProjectionCovered) {
@@ -261,7 +259,6 @@ TEST(ParsedProjectionTest, SortKeyMetaProjectionCovered) {
ASSERT_FALSE(parsedProjection->requiresMatchDetails());
ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
- ASSERT_FALSE(parsedProjection->wantIndexKey());
}
TEST(ParsedProjectionTest, SortKeyMetaAndSlice) {
@@ -276,7 +273,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndSlice) {
ASSERT_FALSE(parsedProjection->requiresMatchDetails());
ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
- ASSERT_FALSE(parsedProjection->wantIndexKey());
}
TEST(ParsedProjectionTest, SortKeyMetaAndElemMatch) {
@@ -291,7 +287,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndElemMatch) {
ASSERT_FALSE(parsedProjection->requiresMatchDetails());
ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
- ASSERT_FALSE(parsedProjection->wantIndexKey());
}
TEST(ParsedProjectionTest, SortKeyMetaAndExclusion) {
@@ -305,7 +300,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndExclusion) {
ASSERT_FALSE(parsedProjection->requiresMatchDetails());
ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
- ASSERT_FALSE(parsedProjection->wantIndexKey());
}
//
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 3487e955675..0d09172ad37 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -306,8 +306,8 @@ auto isCoveredOrAlreadyFetched(const vector<StringData>& fields,
* Checks all properties that exclude a projection from being simple.
*/
auto isSimpleProjection(const CanonicalQuery& query) {
- return !query.getProj()->wantIndexKey() && !query.getProj()->wantSortKey() &&
- !query.getProj()->hasDottedFieldPath() && !query.getProj()->requiresDocument();
+ return !query.getProj()->wantSortKey() && !query.getProj()->hasDottedFieldPath() &&
+ !query.getProj()->requiresDocument();
}
/**
@@ -332,6 +332,21 @@ auto produceCoveredKeyObj(QuerySolutionNode* solnRoot) {
}
/**
+ * Adds a stage to generate the sort key metadata if there's no sort stage but we have a sortKey
+ * meta-projection.
+ */
+std::unique_ptr<QuerySolutionNode> addSortKeyGeneratorStageIfNeeded(
+ const CanonicalQuery& query, bool hasSortStage, std::unique_ptr<QuerySolutionNode> solnRoot) {
+ if (!hasSortStage && query.getProj() && query.getProj()->wantSortKey()) {
+ auto keyGenNode = std::make_unique<SortKeyGeneratorNode>();
+ keyGenNode->sortSpec = query.getQueryRequest().getSort();
+ keyGenNode->children.push_back(solnRoot.release());
+ return keyGenNode;
+ }
+ return solnRoot;
+}
+
+/**
* When projection needs to be added to the solution tree, this function chooses between the default
* implementation and one of the fast paths.
*/
@@ -340,24 +355,12 @@ std::unique_ptr<ProjectionNode> analyzeProjection(const CanonicalQuery& query,
const bool hasSortStage) {
const QueryRequest& qr = query.getQueryRequest();
- // If there's no sort stage but we have a sortKey meta-projection, we need to add a stage to
- // generate the sort key metadata.
- auto addSortKeyGeneratorStageIfNeeded = [&]() {
- if (!hasSortStage && query.getProj()->wantSortKey()) {
- auto keyGenNode = std::make_unique<SortKeyGeneratorNode>();
- keyGenNode->sortSpec = qr.getSort();
- keyGenNode->children.push_back(solnRoot.release());
- solnRoot = std::move(keyGenNode);
- }
- };
-
LOG(5) << "PROJECTION: Current plan is:\n" << redact(solnRoot->toString());
// If the projection requires the entire document we add a fetch stage if not present. Otherwise
- // we add a fetch stage if we are not covered and not returnKey.
+ // we add a fetch stage if we are not covered.
if ((query.getProj()->requiresDocument() && !solnRoot->fetched()) ||
- (!isCoveredOrAlreadyFetched(query.getProj()->getRequiredFields(), *solnRoot) &&
- !query.getProj()->wantIndexKey())) {
+ (!isCoveredOrAlreadyFetched(query.getProj()->getRequiredFields(), *solnRoot))) {
auto fetch = std::make_unique<FetchNode>();
fetch->children.push_back(solnRoot.release());
solnRoot = std::move(fetch);
@@ -372,30 +375,33 @@ std::unique_ptr<ProjectionNode> analyzeProjection(const CanonicalQuery& query,
if (isSimpleProjection(query)) {
// If the projection is simple, but not covered, use 'ProjectionNodeSimple'.
if (solnRoot->fetched()) {
- addSortKeyGeneratorStageIfNeeded();
return std::make_unique<ProjectionNodeSimple>(
- std::move(solnRoot), *query.root(), qr.getProj(), *query.getProj());
+ addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+ *query.root(),
+ qr.getProj(),
+ *query.getProj());
} else {
// If we're here we're not fetched so we're covered. Let's see if we can get out of
// using the default projType. If 'solnRoot' is an index scan we can use the faster
// covered impl.
BSONObj coveredKeyObj = produceCoveredKeyObj(solnRoot.get());
if (!coveredKeyObj.isEmpty()) {
- addSortKeyGeneratorStageIfNeeded();
- return std::make_unique<ProjectionNodeCovered>(std::move(solnRoot),
- *query.root(),
- qr.getProj(),
- *query.getProj(),
- std::move(coveredKeyObj));
+ return std::make_unique<ProjectionNodeCovered>(
+ addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+ *query.root(),
+ qr.getProj(),
+ *query.getProj(),
+ std::move(coveredKeyObj));
}
}
}
- addSortKeyGeneratorStageIfNeeded();
return std::make_unique<ProjectionNodeDefault>(
- std::move(solnRoot), *query.root(), qr.getProj(), *query.getProj());
+ addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+ *query.root(),
+ qr.getProj(),
+ *query.getProj());
}
-
} // namespace
// static
@@ -797,7 +803,13 @@ std::unique_ptr<QuerySolution> QueryPlannerAnalysis::analyzeDataAccess(
}
// Project the results.
- if (query.getProj()) {
+ if (qr.returnKey()) {
+ // We don't need a projection stage if returnKey was requested since the intended behavior
+ // is that the projection is ignored when returnKey is specified.
+ solnRoot = std::make_unique<ReturnKeyNode>(
+ addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)),
+ QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(qr.getProj()));
+ } else if (query.getProj()) {
solnRoot = analyzeProjection(query, std::move(solnRoot), hasSortStage);
// If we don't have a covered project, and we're not allowed to put an uncovered one in,
// bail out.
diff --git a/src/mongo/db/query/query_planner_common.cpp b/src/mongo/db/query/query_planner_common.cpp
index 35cb1ac0293..2b2117676a8 100644
--- a/src/mongo/db/query/query_planner_common.cpp
+++ b/src/mongo/db/query/query_planner_common.cpp
@@ -75,4 +75,32 @@ void QueryPlannerCommon::reverseScans(QuerySolutionNode* node) {
}
}
+// TODO SERVER-42422: reimplement by walking a projection AST rather than raw bson.
+std::vector<std::string> QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(
+ const BSONObj& proj) {
+ std::vector<std::string> sortKeyMetaFields;
+
+ for (auto&& elem : proj) {
+ if (elem.type() == BSONType::Object) {
+ BSONObj obj = elem.embeddedObject();
+ // The caller must have already validated the projection, so we expect
+ // to see only one element.
+ invariant(1 == obj.nFields());
+
+ BSONElement firstSubElem = obj.firstElement();
+ if (firstSubElem.fieldNameStringData() == "$meta") {
+ invariant(firstSubElem.type() == BSONType::String);
+ if (firstSubElem.valueStringData() == QueryRequest::metaSortKey) {
+ invariant(std::find(sortKeyMetaFields.begin(),
+ sortKeyMetaFields.end(),
+ elem.fieldName()) == sortKeyMetaFields.end());
+ sortKeyMetaFields.push_back(elem.fieldName());
+ }
+ }
+ }
+ }
+
+ return sortKeyMetaFields;
+}
+
} // namespace mongo
diff --git a/src/mongo/db/query/query_planner_common.h b/src/mongo/db/query/query_planner_common.h
index c232291e0e3..76c70742a3c 100644
--- a/src/mongo/db/query/query_planner_common.h
+++ b/src/mongo/db/query/query_planner_common.h
@@ -82,6 +82,14 @@ public:
* the scan direction and index bounds.
*/
static void reverseScans(QuerySolutionNode* node);
+
+ /**
+ * Extracts all field names for the sortKey meta-projection and stores them in the returned
+ * array. Returns an empty array if there were no sortKey meta-projection specified in the
+ * given projection 'proj'. For example, given a projection {a:1, b: {$meta: "sortKey"},
+ * c: {$meta: "sortKey"}}, the returned vector will contain two elements ["b", "c"].
+ */
+ static std::vector<std::string> extractSortKeyMetaFieldsFromProjection(const BSONObj& proj);
};
} // namespace mongo
diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index b3c87b40ab8..fde6441559b 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -58,7 +58,6 @@ const char QueryRequest::queryOptionMaxTimeMS[] = "$maxTimeMS";
const string QueryRequest::metaGeoNearDistance("geoNearDistance");
const string QueryRequest::metaGeoNearPoint("geoNearPoint");
-const string QueryRequest::metaIndexKey("indexKey");
const string QueryRequest::metaRecordId("recordId");
const string QueryRequest::metaSortKey("sortKey");
const string QueryRequest::metaTextScore("textScore");
@@ -592,16 +591,6 @@ void QueryRequest::asFindCommandInternal(BSONObjBuilder* cmdBuilder) const {
}
}
-void QueryRequest::addReturnKeyMetaProj() {
- BSONObjBuilder projBob;
- projBob.appendElements(_proj);
- // We use $$ because it's never going to show up in a user's projection.
- // The exact text doesn't matter.
- BSONObj indexKey = BSON("$$" << BSON("$meta" << QueryRequest::metaIndexKey));
- projBob.append(indexKey.firstElement());
- _proj = projBob.obj();
-}
-
void QueryRequest::addShowRecordIdMetaProj() {
BSONObjBuilder projBob;
projBob.appendElements(_proj);
@@ -937,7 +926,6 @@ Status QueryRequest::initFullQuery(const BSONObj& top) {
// Won't throw.
if (e.trueValue()) {
_returnKey = true;
- addReturnKeyMetaProj();
}
} else if (name == "showDiskLoc") {
// Won't throw.
@@ -1004,11 +992,6 @@ void QueryRequest::initFromInt(int options) {
}
void QueryRequest::addMetaProjection() {
- // We might need to update the projection object with a $meta projection.
- if (returnKey()) {
- addReturnKeyMetaProj();
- }
-
if (showRecordId()) {
addShowRecordIdMetaProj();
}
diff --git a/src/mongo/db/query/query_request.h b/src/mongo/db/query/query_request.h
index 7b0cb74dd80..17e915ace5c 100644
--- a/src/mongo/db/query/query_request.h
+++ b/src/mongo/db/query/query_request.h
@@ -135,7 +135,6 @@ public:
// Names of the $meta projection values.
static const std::string metaGeoNearDistance;
static const std::string metaGeoNearPoint;
- static const std::string metaIndexKey;
static const std::string metaRecordId;
static const std::string metaSortKey;
static const std::string metaTextScore;
@@ -455,11 +454,6 @@ private:
Status initFullQuery(const BSONObj& top);
/**
- * Updates the projection object with a $meta projection for the returnKey option.
- */
- void addReturnKeyMetaProj();
-
- /**
* Updates the projection object with a $meta projection for the showRecordId option.
*/
void addShowRecordIdMetaProj();
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index ddbe87074a8..6b076d80e96 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -31,6 +31,8 @@
#include "mongo/db/query/query_solution.h"
+#include <boost/algorithm/string/join.hpp>
+
#include "mongo/bson/bsontypes.h"
#include "mongo/bson/mutable/document.h"
#include "mongo/bson/simple_bsonelement_comparator.h"
@@ -848,6 +850,27 @@ bool IndexScanNode::operator==(const IndexScanNode& other) const {
}
//
+// ReturnKeyNode
+//
+
+void ReturnKeyNode::appendToString(str::stream* ss, int indent) const {
+ addIndent(ss, indent);
+ *ss << "RETURN_KEY\n";
+ addIndent(ss, indent + 1);
+ *ss << "sortKeyMetaFields = [" << boost::algorithm::join(sortKeyMetaFields, ", ") << "]\n";
+ addCommon(ss, indent);
+ addIndent(ss, indent + 1);
+ *ss << "Child:" << '\n';
+ children[0]->appendToString(ss, indent + 2);
+}
+
+QuerySolutionNode* ReturnKeyNode::clone() const {
+ auto copy = std::make_unique<ReturnKeyNode>(
+ std::unique_ptr<QuerySolutionNode>(children[0]->clone()), std::vector(sortKeyMetaFields));
+ return copy.release();
+}
+
+//
// ProjectionNode
//
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index cba57fd1263..48a0288a317 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -540,6 +540,35 @@ struct IndexScanNode : public QuerySolutionNode {
std::set<StringData> multikeyFields;
};
+struct ReturnKeyNode : public QuerySolutionNode {
+ ReturnKeyNode(std::unique_ptr<QuerySolutionNode> child,
+ std::vector<std::string> sortKeyMetaFields)
+ : QuerySolutionNode(std::move(child)), sortKeyMetaFields(std::move(sortKeyMetaFields)) {}
+
+ StageType getType() const final {
+ return STAGE_RETURN_KEY;
+ }
+
+ void appendToString(str::stream* ss, int indent) const final;
+
+ bool fetched() const final {
+ return children[0]->fetched();
+ }
+ bool hasField(const std::string& field) const final {
+ return false;
+ }
+ bool sortedByDiskLoc() const final {
+ return children[0]->sortedByDiskLoc();
+ }
+ const BSONObjSet& getSort() const final {
+ return children[0]->getSort();
+ }
+
+ QuerySolutionNode* clone() const final;
+
+ std::vector<std::string> sortKeyMetaFields;
+};
+
/**
* We have a few implementations of the projection functionality. They are chosen by constructing
* a type derived from this abstract struct. The most general implementation 'ProjectionNodeDefault'
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);
diff --git a/src/mongo/db/query/stage_types.h b/src/mongo/db/query/stage_types.h
index 8b34d4dfecc..17a67f1169f 100644
--- a/src/mongo/db/query/stage_types.h
+++ b/src/mongo/db/query/stage_types.h
@@ -87,6 +87,7 @@ enum StageType {
STAGE_QUEUED_DATA,
STAGE_RECORD_STORE_FAST_COUNT,
+ STAGE_RETURN_KEY,
STAGE_SHARDING_FILTER,
STAGE_SKIP,
STAGE_SORT,
View Lorry Controller Status