SERVER-42423 Use ProjectionExecutor in ProjectionStageDefault

6 files changed, 128 insertions, 93 deletions

diff --git a/src/mongo/db/exec/find_projection_executor.cpp b/src/mongo/db/exec/find_projection_executor.cpp
index 098cdb93d71..9ff1bcd4602 100644
--- a/src/mongo/db/exec/find_projection_executor.cpp
+++ b/src/mongo/db/exec/find_projection_executor.cpp
@@ -128,17 +128,19 @@ Value applySliceProjectionHelper(const Document& input,
                                  size_t fieldPathIndex) {
     invariant(fieldPathIndex < params.path.getPathLength());
 
-    auto fieldName = params.path.getFieldName(fieldPathIndex);
+    auto fieldName = params.path.getFieldName(fieldPathIndex++);
     Value val{input[fieldName]};
 
     switch (val.getType()) {
         case BSONType::Array:
-            val = (fieldPathIndex + 1 == params.path.getPathLength())
+            val = (fieldPathIndex == params.path.getPathLength())
                 ? sliceArray(val.getArray(), params.skip, params.limit)
-                : applySliceProjectionToArray(val.getArray(), params, fieldPathIndex + 1);
+                : applySliceProjectionToArray(val.getArray(), params, fieldPathIndex);
             break;
         case BSONType::Object:
-            val = applySliceProjectionHelper(val.getDocument(), params, fieldPathIndex + 1);
+            if (fieldPathIndex < params.path.getPathLength()) {
+                val = applySliceProjectionHelper(val.getDocument(), params, fieldPathIndex);
+            }
             break;
         default:
             break;
diff --git a/src/mongo/db/exec/find_projection_executor_test.cpp b/src/mongo/db/exec/find_projection_executor_test.cpp
index 2bff975ed8c..39194fbdd93 100644
--- a/src/mongo/db/exec/find_projection_executor_test.cpp
+++ b/src/mongo/db/exec/find_projection_executor_test.cpp
@@ -270,6 +270,10 @@ TEST(SliceProjection, CorrectlyProjectsSimplePath) {
     doc = Document{fromjson("{a: 1}")};
     ASSERT_DOCUMENT_EQ(Document{fromjson("{a: 1}")},
                        applySliceProjection(doc, "a", boost::none, 2));
+
+    doc = Document{fromjson("{a: {b: 1}}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: {b: 1}}")},
+                       applySliceProjection(doc, "a", boost::none, 2));
 }
 
 TEST(SliceProjection, CorrectlyProjectsDottedPath) {
@@ -293,6 +297,10 @@ TEST(SliceProjection, CorrectlyProjectsDottedPath) {
     ASSERT_DOCUMENT_EQ(Document{fromjson("{a: {b: 1}}")},
                        applySliceProjection(doc, "a.b", boost::none, 2));
 
+    doc = Document{fromjson("{a: {b: {c: 1}}}")};
+    ASSERT_DOCUMENT_EQ(Document{fromjson("{a: {b: {c: 1}}}")},
+                       applySliceProjection(doc, "a.b", boost::none, 2));
+
     doc = Document{fromjson("{a: [{b: [1,2,3], c: 1}]}")};
     ASSERT_DOCUMENT_EQ(Document{fromjson("{a: [{b: [3], c: 1}]}")},
                        applySliceProjection(doc, "a.b", boost::none, -1));
diff --git a/src/mongo/db/exec/projection.cpp b/src/mongo/db/exec/projection.cpp
index f334655562e..5d50aee547c 100644
--- a/src/mongo/db/exec/projection.cpp
+++ b/src/mongo/db/exec/projection.cpp
@@ -50,68 +50,68 @@ static const char* kIdField = "_id";
 
 namespace {
 
-BSONObj indexKey(const WorkingSetMember& member) {
-    return member.metadata().getIndexKey();
+void transitionMemberToOwnedObj(Document&& doc, WorkingSetMember* member) {
+    member->keyData.clear();
+    member->recordId = {};
+    member->doc = {{}, std::move(doc)};
+    member->transitionToOwnedObj();
 }
 
-BSONObj sortKey(const WorkingSetMember& member) {
-    return DocumentMetadataFields::serializeSortKey(member.metadata().isSingleElementKey(),
-                                                    member.metadata().getSortKey());
+void transitionMemberToOwnedObj(const BSONObj& bo, WorkingSetMember* member) {
+    transitionMemberToOwnedObj(Document{bo}, member);
 }
 
-double geoDistance(const WorkingSetMember& member) {
-    return member.metadata().getGeoNearDistance();
+/**
+ * Moves document metadata fields from the WSM into the given document 'doc', and returns the same
+ * document but with populated metadata.
+ */
+auto attachMetadataToDocument(Document&& doc, WorkingSetMember* member) {
+    MutableDocument md{std::move(doc)};
+    md.setMetadata(member->releaseMetadata());
+    return md.freeze();
 }
 
-Value geoPoint(const WorkingSetMember& member) {
-    return member.metadata().getGeoNearPoint();
+/**
+ * Moves document metadata fields from the document 'doc' into the WSM, and returns the same
+ * document but without metadata.
+ */
+auto attachMetadataToWorkingSetMember(Document&& doc, WorkingSetMember* member) {
+    MutableDocument md{std::move(doc)};
+    member->setMetadata(md.releaseMetadata());
+    return md.freeze();
 }
 
-double textScore(const WorkingSetMember& member) {
-    auto&& metadata = member.metadata();
-    if (metadata.hasTextScore()) {
-        return metadata.getTextScore();
-    } else {
-        // It is permitted to request a text score when none has been computed. Zero is returned as
-        // an empty value in this case.
-        return 0.0;
+/**
+ * Given an index key 'dehyratedKey' with no field names, returns a new Document representing the
+ * index key after adding field names according to 'keyPattern'.
+ *
+ * For example, given:
+ *    - the 'keyPatern' of {'a.b': 1, c: 1}
+ *    - the 'dehydratedKey' of {'': 'abc', '': 10}
+ *
+ * The resulting document will be: {a: {b: 'abc'}, c: 10}
+ */
+auto rehydrateIndexKey(const BSONObj& keyPattern, const BSONObj& dehydratedKey) {
+    MutableDocument md;
+    BSONObjIterator keyIter{keyPattern};
+    BSONObjIterator valueIter{dehydratedKey};
+
+    while (keyIter.more() && valueIter.more()) {
+        auto fieldName = keyIter.next().fieldNameStringData();
+        auto value = valueIter.next();
+
+        // Skip the $** index virtual field, as it's not part of the actual index key.
+        if (fieldName == "$_path") {
+            continue;
+        }
+
+        md.setNestedField(fieldName, Value{value});
     }
-}
 
-void transitionMemberToOwnedObj(const BSONObj& bo, WorkingSetMember* member) {
-    member->keyData.clear();
-    member->recordId = RecordId();
-    member->resetDocument(SnapshotId(), bo);
-    member->transitionToOwnedObj();
-}
+    invariant(!keyIter.more());
+    invariant(!valueIter.more());
 
-StatusWith<BSONObj> provideMetaFieldsAndPerformExec(const ProjectionExec& exec,
-                                                    const WorkingSetMember& member) {
-    if (exec.needsGeoNearDistance() && !member.metadata().hasGeoNearDistance())
-        return Status(ErrorCodes::InternalError, "near loc dist requested but no data available");
-
-    if (exec.needsGeoNearPoint() && !member.metadata().hasGeoNearPoint())
-        return Status(ErrorCodes::InternalError, "near loc proj requested but no data available");
-
-    return member.hasObj()
-        ? exec.project(member.doc.value().toBson(),
-                       exec.needsGeoNearDistance()
-                           ? boost::optional<const double>(geoDistance(member))
-                           : boost::none,
-                       exec.needsGeoNearPoint() ? geoPoint(member) : Value{},
-                       exec.needsSortKey() ? sortKey(member) : BSONObj(),
-                       exec.needsTextScore() ? boost::optional<const double>(textScore(member))
-                                             : boost::none,
-                       member.recordId.repr())
-        : exec.projectCovered(
-              member.keyData,
-              exec.needsGeoNearDistance() ? boost::optional<const double>(geoDistance(member))
-                                          : boost::none,
-              exec.needsGeoNearPoint() ? geoPoint(member) : Value{},
-              exec.needsSortKey() ? sortKey(member) : BSONObj(),
-              exec.needsTextScore() ? boost::optional<const double>(textScore(member))
-                                    : boost::none,
-              member.recordId.repr());
+    return md.freeze();
 }
 }  // namespace
 
@@ -191,26 +191,52 @@ std::unique_ptr<PlanStageStats> ProjectionStage::getStats() {
     return ret;
 }
 
-ProjectionStageDefault::ProjectionStageDefault(OperationContext* opCtx,
+ProjectionStageDefault::ProjectionStageDefault(boost::intrusive_ptr<ExpressionContext> expCtx,
                                                const BSONObj& projObj,
+                                               const projection_ast::Projection* projection,
                                                WorkingSet* ws,
-                                               std::unique_ptr<PlanStage> child,
-                                               const MatchExpression& fullExpression,
-                                               const CollatorInterface* collator)
-    : ProjectionStage(opCtx, projObj, ws, std::move(child), "PROJECTION_DEFAULT"),
-      _exec(opCtx, projObj, &fullExpression, collator) {}
+                                               std::unique_ptr<PlanStage> child)
+    : ProjectionStage{expCtx->opCtx, projObj, ws, std::move(child), "PROJECTION_DEFAULT"},
+      _wantRecordId{projection->metadataDeps()[DocumentMetadataFields::kRecordId]},
+      _projectType{projection->type()},
+      _executor{projection_executor::buildProjectionExecutor(expCtx, projection, {})} {}
 
 Status ProjectionStageDefault::transform(WorkingSetMember* member) const {
-    // The default no-fast-path case.
-    if (_exec.needsSortKey() && !member->metadata().hasSortKey())
-        return Status(ErrorCodes::InternalError,
-                      "sortKey meta-projection requested but no data available");
+    Document input;
+
+    // Most metadata should have already been stored within the WSM when we project out a document.
+    // The recordId metadata is different though, because it's a fundamental part of the WSM and
+    // we store it within the WSM itself rather than WSM metadata, so we need to transfer it into
+    // the metadata object if the projection has a recordId $meta expression.
+    if (_wantRecordId && !member->metadata().hasRecordId()) {
+        member->metadata().setRecordId(member->recordId);
+    }
 
-    auto projected = provideMetaFieldsAndPerformExec(_exec, *member);
-    if (!projected.isOK())
-        return projected.getStatus();
+    if (member->hasObj()) {
+        input = std::move(member->doc.value());
+    } else {
+        // We have a covered projection, which is only supported in inclusion mode.
+        invariant(_projectType == projection_ast::ProjectType::kInclusion);
+        // We're pulling data from an index key, so there must be exactly one key entry in the WSM
+        // as the planner guarantees that it will never generate a covered plan in the case of index
+        // intersection.
+        invariant(member->keyData.size() == 1);
+
+        // For covered projection we will rehydrate in index key into a Document and then pass it
+        // through the projection executor to include only required fields, including metadata
+        // fields.
+        input = rehydrateIndexKey(member->keyData[0].indexKeyPattern, member->keyData[0].keyData);
+    }
 
-    transitionMemberToOwnedObj(projected.getValue(), member);
+    // Before applying the projection we will move document metadata from the WSM into the document
+    // itself, in case the projection contains $meta expressions and needs this data, and will move
+    // it back to the WSM once the projection has been applied.
+    auto projected = attachMetadataToWorkingSetMember(
+        _executor->applyTransformation(attachMetadataToDocument(std::move(input), member)), member);
+    // An exclusion projection can return an unowned object since the output document is
+    // constructed from the input one backed by BSON which is owned by the storage system, so we
+    // need to  make sure we transition an owned document.
+    transitionMemberToOwnedObj(projected.getOwned(), member);
 
     return Status::OK();
 }
diff --git a/src/mongo/db/exec/projection.h b/src/mongo/db/exec/projection.h
index e4ad2ef8af6..5b368750cad 100644
--- a/src/mongo/db/exec/projection.h
+++ b/src/mongo/db/exec/projection.h
@@ -30,15 +30,13 @@
 #pragma once
 
 #include "mongo/db/exec/plan_stage.h"
-#include "mongo/db/exec/projection_exec.h"
+#include "mongo/db/exec/projection_executor.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/matcher/expression.h"
+#include "mongo/db/query/projection_ast.h"
 #include "mongo/db/record_id.h"
 
 namespace mongo {
-
-class CollatorInterface;
-
 /**
  * This stage computes a projection. This is an abstract base class for various projection
  * implementations.
@@ -101,12 +99,11 @@ public:
     /**
      * ProjectionNodeDefault should use this for construction.
      */
-    ProjectionStageDefault(OperationContext* opCtx,
+    ProjectionStageDefault(boost::intrusive_ptr<ExpressionContext> expCtx,
                            const BSONObj& projObj,
+                           const projection_ast::Projection* projection,
                            WorkingSet* ws,
-                           std::unique_ptr<PlanStage> child,
-                           const MatchExpression& fullExpression,
-                           const CollatorInterface* collator);
+                           std::unique_ptr<PlanStage> child);
 
     StageType stageType() const final {
         return STAGE_PROJECTION_DEFAULT;
@@ -115,8 +112,10 @@ public:
 private:
     Status transform(WorkingSetMember* member) const final;
 
-    // Fully-general heavy execution object.
-    ProjectionExec _exec;
+    // True, if the projection contains a recordId $meta expression.
+    const bool _wantRecordId;
+    const projection_ast::ProjectType _projectType;
+    std::unique_ptr<parsed_aggregation_projection::ParsedAggregationProjection> _executor;
 };
 
 /**
diff --git a/src/mongo/db/exec/projection_executor.cpp b/src/mongo/db/exec/projection_executor.cpp
index 91c56032cd6..29192f19547 100644
--- a/src/mongo/db/exec/projection_executor.cpp
+++ b/src/mongo/db/exec/projection_executor.cpp
@@ -97,7 +97,7 @@ auto makeProjectionPostImageExpression(const ProjectionExecutorVisitorData<Execu
  * a field path expressions will be created to access a projection post-image document.
  */
 template <typename Executor>
-auto createFindPositionalExpression(projection_ast::ProjectionPositionalASTNode* node,
+auto createFindPositionalExpression(const projection_ast::ProjectionPositionalASTNode* node,
                                     const ProjectionExecutorVisitorData<Executor>& data,
                                     const FieldPath& path) {
     invariant(node);
@@ -123,7 +123,7 @@ auto createFindPositionalExpression(projection_ast::ProjectionPositionalASTNode*
  * a field path expressions will be created to access a projection post-image document.
  */
 template <typename Executor>
-auto createFindSliceExpression(projection_ast::ProjectionSliceASTNode* node,
+auto createFindSliceExpression(const projection_ast::ProjectionSliceASTNode* node,
                                const ProjectionExecutorVisitorData<Executor>& data,
                                const FieldPath& path) {
     invariant(node);
@@ -137,7 +137,7 @@ auto createFindSliceExpression(projection_ast::ProjectionSliceASTNode* node,
  * 'path' on the input document.
  */
 template <typename Executor>
-auto createFindElemMatchExpression(projection_ast::ProjectionElemMatchASTNode* node,
+auto createFindElemMatchExpression(const projection_ast::ProjectionElemMatchASTNode* node,
                                    const ProjectionExecutorVisitorData<Executor>& data,
                                    const FieldPath& path) {
     invariant(node);
@@ -166,14 +166,14 @@ auto createFindElemMatchExpression(projection_ast::ProjectionElemMatchASTNode* n
  * 'PathTrackingVisitorContext'.
  */
 template <typename Executor>
-class ProjectionExecutorVisitor final : public projection_ast::ProjectionASTVisitor {
+class ProjectionExecutorVisitor final : public projection_ast::ProjectionASTConstVisitor {
 public:
     ProjectionExecutorVisitor(ProjectionExecutorVisitorContext<Executor>* context)
         : _context{context} {
         invariant(_context);
     }
 
-    void visit(projection_ast::ProjectionPositionalASTNode* node) final {
+    void visit(const projection_ast::ProjectionPositionalASTNode* node) final {
         constexpr auto isInclusion = std::is_same_v<Executor, ParsedInclusionProjection>;
         invariant(isInclusion);
 
@@ -184,7 +184,7 @@ public:
         userData.setRootReplacementExpression(createFindPositionalExpression(node, userData, path));
     }
 
-    void visit(projection_ast::ProjectionSliceASTNode* node) final {
+    void visit(const projection_ast::ProjectionSliceASTNode* node) final {
         const auto& path = _context->fullPath();
         auto& userData = _context->data();
 
@@ -198,7 +198,7 @@ public:
         userData.setRootReplacementExpression(createFindSliceExpression(node, userData, path));
     }
 
-    void visit(projection_ast::ProjectionElemMatchASTNode* node) final {
+    void visit(const projection_ast::ProjectionElemMatchASTNode* node) final {
         const auto& path = _context->fullPath();
         const auto& userData = _context->data();
 
@@ -206,14 +206,14 @@ public:
             path.fullPath(), createFindElemMatchExpression(node, userData, path));
     }
 
-    void visit(projection_ast::ExpressionASTNode* node) final {
+    void visit(const projection_ast::ExpressionASTNode* node) final {
         const auto& path = _context->fullPath();
         const auto& userData = _context->data();
 
         userData.rootNode()->addExpressionForPath(path.fullPath(), node->expression());
     }
 
-    void visit(projection_ast::BooleanConstantASTNode* node) final {
+    void visit(const projection_ast::BooleanConstantASTNode* node) final {
         const auto& path = _context->fullPath();
         const auto& userData = _context->data();
 
@@ -231,8 +231,8 @@ public:
         userData.rootNode()->addProjectionForPath(path.fullPath());
     }
 
-    void visit(projection_ast::ProjectionPathASTNode* node) final {}
-    void visit(projection_ast::MatchExpressionASTNode* node) final {}
+    void visit(const projection_ast::ProjectionPathASTNode* node) final {}
+    void visit(const projection_ast::MatchExpressionASTNode* node) final {}
 
 private:
     ProjectionExecutorVisitorContext<Executor>* _context;
@@ -245,7 +245,7 @@ private:
  */
 template <typename Executor>
 auto buildProjectionExecutor(boost::intrusive_ptr<ExpressionContext> expCtx,
-                             projection_ast::ProjectionPathASTNode* root,
+                             const projection_ast::ProjectionPathASTNode* root,
                              const ProjectionPolicies policies) {
     ProjectionExecutorVisitorContext<Executor> context{
         {std::make_unique<Executor>(expCtx, policies), expCtx}};
@@ -258,7 +258,7 @@ auto buildProjectionExecutor(boost::intrusive_ptr<ExpressionContext> expCtx,
 
 std::unique_ptr<ParsedAggregationProjection> buildProjectionExecutor(
     boost::intrusive_ptr<ExpressionContext> expCtx,
-    projection_ast::Projection* projection,
+    const projection_ast::Projection* projection,
     const ProjectionPolicies policies) {
     invariant(projection);
 
diff --git a/src/mongo/db/exec/projection_executor.h b/src/mongo/db/exec/projection_executor.h
index 98e7fe89bfd..e3b7b2b24dc 100644
--- a/src/mongo/db/exec/projection_executor.h
+++ b/src/mongo/db/exec/projection_executor.h
@@ -40,6 +40,6 @@ namespace mongo::projection_executor {
  */
 std::unique_ptr<parsed_aggregation_projection::ParsedAggregationProjection> buildProjectionExecutor(
     boost::intrusive_ptr<ExpressionContext> expCtx,
-    projection_ast::Projection* projection,
+    const projection_ast::Projection* projection,
     ProjectionPolicies policies);
 }  // namespace mongo::projection_executor