SERVER-43830 use ProjectionAST in when extracting meta sort key for returnKey

11 files changed, 97 insertions, 43 deletions

diff --git a/jstests/core/return_key.js b/jstests/core/return_key.js
index 684b1952310..771184cccef 100644
--- a/jstests/core/return_key.js
+++ b/jstests/core/return_key.js
@@ -81,4 +81,12 @@ assert.eq(results, [
     {a: 2, c: {'': 2}, d: {'': 2}},
     {a: 1, c: {'': 3}, d: {'': 3}}
 ]);
+
+// returnKey with a sortKey $meta projection on a nested field.
+// TODO: SERVER-42435: Re-enable this test once 4.2 find() projection parser is removed.
+if (false) {
+    results =
+        coll.find({}, {"c.d": {$meta: "sortKey"}}).hint({a: 1}).sort({b: 1}).returnKey().toArray();
+    assert.eq(results, [{a: 3, c: {d: {'': 1}}}, {a: 2, c: {d: {'': 1}}}, {a: 1, c: {d: {'': 1}}}]);
+}
 })();
diff --git a/src/mongo/db/exec/return_key.cpp b/src/mongo/db/exec/return_key.cpp
index 7b8bf6cbf77..c133d0418f2 100644
--- a/src/mongo/db/exec/return_key.cpp
+++ b/src/mongo/db/exec/return_key.cpp
@@ -33,7 +33,9 @@
 
 #include "mongo/db/exec/return_key.h"
 
+#include "mongo/db/exec/document_value/document.h"
 #include "mongo/db/exec/working_set_common.h"
+#include "mongo/db/pipeline/field_path.h"
 #include "mongo/util/log.h"
 
 namespace mongo {
@@ -86,19 +88,20 @@ Status ReturnKeyStage::_extractIndexKey(WorkingSetMember* member) {
         ? DocumentMetadataFields::serializeSortKey(member->metadata().isSingleElementKey(),
                                                    member->metadata().getSortKey())
         : BSONObj();
-    BSONObjBuilder bob;
 
-    if (!indexKey.isEmpty()) {
-        bob.appendElements(indexKey);
+    MutableDocument md;
+
+    for (auto&& elem : indexKey) {
+        md.addField(elem.fieldNameStringData(), Value(elem));
     }
 
-    for (const auto& fieldName : _sortKeyMetaFields) {
-        bob.append(fieldName, sortKey);
+    for (const auto& fieldPath : _sortKeyMetaFields) {
+        md.setNestedField(fieldPath, Value(sortKey));
     }
 
     member->keyData.clear();
     member->recordId = {};
-    member->resetDocument({}, bob.obj());
+    member->doc = {{}, md.freeze()};
     member->transitionToOwnedObj();
 
     return Status::OK();
diff --git a/src/mongo/db/exec/return_key.h b/src/mongo/db/exec/return_key.h
index 7212de12843..9a6f2926b2d 100644
--- a/src/mongo/db/exec/return_key.h
+++ b/src/mongo/db/exec/return_key.h
@@ -47,7 +47,7 @@ public:
     static constexpr StringData kStageName = "RETURN_KEY"_sd;
 
     ReturnKeyStage(OperationContext* opCtx,
-                   std::vector<std::string> sortKeyMetaFields,
+                   std::vector<FieldPath> sortKeyMetaFields,
                    WorkingSet* ws,
                    std::unique_ptr<PlanStage> child)
         : PlanStage(opCtx, std::move(child), kStageName.rawData()),
@@ -78,6 +78,6 @@ private:
 
     // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey
     // meta-projection.
-    std::vector<std::string> _sortKeyMetaFields;
+    std::vector<FieldPath> _sortKeyMetaFields;
 };
 }  // namespace mongo
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index dec282568dc..ced708e90dc 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -390,6 +390,8 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx,
                 std::move(root));
         }
 
+        const auto* cqProjection = canonicalQuery->getProj();
+
         // Add a SortKeyGeneratorStage if the query requested sortKey metadata.
         if (canonicalQuery->metadataDeps()[DocumentMetadataFields::kSortKey]) {
             root = std::make_unique<SortKeyGeneratorStage>(
@@ -406,16 +408,17 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx,
             // returnKey.
             root = std::make_unique<ReturnKeyStage>(
                 opCtx,
-                QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(
-                    canonicalQuery->getQueryRequest().getProj()),
+                cqProjection
+                    ? QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(*cqProjection)
+                    : std::vector<FieldPath>{},
                 ws,
                 std::move(root));
-        } else if (canonicalQuery->getProj()) {
+        } else if (cqProjection) {
             // 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()->isSimple()) {
+            if (!cqProjection->isSimple()) {
                 root = std::make_unique<ProjectionStageDefault>(
                     canonicalQuery->getExpCtx(),
                     canonicalQuery->getQueryRequest().getProj(),
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 2e6ae124505..27405653ae7 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -795,7 +795,9 @@ std::unique_ptr<QuerySolution> QueryPlannerAnalysis::analyzeDataAccess(
         // 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()));
+            query.getProj()
+                ? QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(*query.getProj())
+                : std::vector<FieldPath>{});
     } 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,
diff --git a/src/mongo/db/query/projection.h b/src/mongo/db/query/projection.h
index 486f9ed9abf..be90f0e8311 100644
--- a/src/mongo/db/query/projection.h
+++ b/src/mongo/db/query/projection.h
@@ -61,11 +61,11 @@ class Projection {
 public:
     Projection(ProjectionPathASTNode root, ProjectType type);
 
-    const ProjectionPathASTNode* root() const {
+    ProjectionPathASTNode* root() {
         return &_root;
     }
 
-    ProjectionPathASTNode* root() {
+    const ProjectionPathASTNode* root() const {
         return &_root;
     }
 
diff --git a/src/mongo/db/query/projection_ast_path_tracking_visitor.h b/src/mongo/db/query/projection_ast_path_tracking_visitor.h
index affc3732dd6..2d989b88ae6 100644
--- a/src/mongo/db/query/projection_ast_path_tracking_visitor.h
+++ b/src/mongo/db/query/projection_ast_path_tracking_visitor.h
@@ -231,5 +231,12 @@ private:
     std::vector<ProjectionASTVisitor<IsConst>*> _preVisitors;
     std::vector<ProjectionASTVisitor<IsConst>*> _postVisitors;
 };
+
+template <class UserData>
+using PathTrackingConstWalker = PathTrackingWalker<UserData, true>;
+
+template <class UserData>
+using PathTrackingMutableWalker = PathTrackingWalker<UserData, false>;
+
 }  // namespace projection_ast
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_planner_common.cpp b/src/mongo/db/query/query_planner_common.cpp
index 2b2117676a8..c55766195bb 100644
--- a/src/mongo/db/query/query_planner_common.cpp
+++ b/src/mongo/db/query/query_planner_common.cpp
@@ -31,6 +31,9 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/base/exact_cast.h"
+#include "mongo/db/query/projection_ast_path_tracking_visitor.h"
+#include "mongo/db/query/projection_ast_walker.h"
 #include "mongo/db/query/query_planner_common.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/log.h"
@@ -75,32 +78,51 @@ 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());
-                }
-            }
+namespace {
+
+struct MetaFieldData {
+    std::vector<FieldPath> metaPaths;
+};
+
+using MetaFieldVisitorContext = projection_ast::PathTrackingVisitorContext<MetaFieldData>;
+
+/**
+ * Visitor which produces a list of paths where $meta expressions are.
+ */
+class MetaFieldVisitor final : public projection_ast::ProjectionASTConstVisitor {
+public:
+    MetaFieldVisitor(MetaFieldVisitorContext* context) : _context(context) {}
+
+
+    void visit(const projection_ast::ExpressionASTNode* node) final {
+        const auto* metaExpr = exact_pointer_cast<const ExpressionMeta*>(node->expressionRaw());
+        if (!metaExpr || metaExpr->getMetaType() != DocumentMetadataFields::MetaType::kSortKey) {
+            return;
         }
+
+        _context->data().metaPaths.push_back(_context->fullPath());
     }
 
-    return sortKeyMetaFields;
-}
+    void visit(const projection_ast::ProjectionPositionalASTNode* node) final {}
+    void visit(const projection_ast::ProjectionSliceASTNode* node) final {}
+    void visit(const projection_ast::ProjectionElemMatchASTNode* node) final {}
+    void visit(const projection_ast::BooleanConstantASTNode* node) final {}
+    void visit(const projection_ast::ProjectionPathASTNode* node) final {}
+    void visit(const projection_ast::MatchExpressionASTNode* node) final {}
+
+private:
+    MetaFieldVisitorContext* _context;
+};
+}  // namespace
 
+std::vector<FieldPath> QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(
+    const projection_ast::Projection& proj) {
+
+    MetaFieldVisitorContext ctx;
+    MetaFieldVisitor visitor(&ctx);
+    projection_ast::PathTrackingConstWalker<MetaFieldData> walker{&ctx, {&visitor}, {}};
+    projection_ast_walker::walk(&walker, proj.root());
+
+    return std::move(ctx.data().metaPaths);
+}
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_planner_common.h b/src/mongo/db/query/query_planner_common.h
index 76c70742a3c..3c3bb88936c 100644
--- a/src/mongo/db/query/query_planner_common.h
+++ b/src/mongo/db/query/query_planner_common.h
@@ -31,6 +31,7 @@
 
 #include "mongo/db/jsobj.h"
 #include "mongo/db/matcher/expression.h"
+#include "mongo/db/query/projection.h"
 #include "mongo/db/query/query_solution.h"
 
 namespace mongo {
@@ -89,7 +90,8 @@ public:
      * 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);
+    static std::vector<FieldPath> extractSortKeyMetaFieldsFromProjection(
+        const projection_ast::Projection& proj);
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index 245621eb388..320f54978a4 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -32,6 +32,7 @@
 #include "mongo/db/query/query_solution.h"
 
 #include <boost/algorithm/string/join.hpp>
+#include <boost/range/adaptor/transformed.hpp>
 
 #include "mongo/bson/bsontypes.h"
 #include "mongo/bson/mutable/document.h"
@@ -858,7 +859,13 @@ 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";
+
+    *ss << "sortKeyMetaFields = ["
+        << boost::algorithm::join(
+               sortKeyMetaFields |
+                   boost::adaptors::transformed([](const auto& field) { return field.fullPath(); }),
+               ", ");
+    *ss << "]\n";
     addCommon(ss, indent);
     addIndent(ss, indent + 1);
     *ss << "Child:" << '\n';
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index dd4eeb20839..109241a2d54 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -542,7 +542,7 @@ struct IndexScanNode : public QuerySolutionNode {
 
 struct ReturnKeyNode : public QuerySolutionNode {
     ReturnKeyNode(std::unique_ptr<QuerySolutionNode> child,
-                  std::vector<std::string> sortKeyMetaFields)
+                  std::vector<FieldPath> sortKeyMetaFields)
         : QuerySolutionNode(std::move(child)), sortKeyMetaFields(std::move(sortKeyMetaFields)) {}
 
     StageType getType() const final {
@@ -566,7 +566,7 @@ struct ReturnKeyNode : public QuerySolutionNode {
 
     QuerySolutionNode* clone() const final;
 
-    std::vector<std::string> sortKeyMetaFields;
+    std::vector<FieldPath> sortKeyMetaFields;
 };
 
 /**