SERVER-63753 Translate $lookup result object creation in SBE

6 files changed, 340 insertions, 258 deletions

diff --git a/src/mongo/db/query/plan_explainer_sbe.cpp b/src/mongo/db/query/plan_explainer_sbe.cpp
index 0f1a472bec3..da13bd87687 100644
--- a/src/mongo/db/query/plan_explainer_sbe.cpp
+++ b/src/mongo/db/query/plan_explainer_sbe.cpp
@@ -175,7 +175,7 @@ void statsToBSON(const QuerySolutionNode* node,
             bob->append("foreignCollection", eln->foreignCollection);
             bob->append("localField", eln->joinFieldLocal);
             bob->append("foreignField", eln->joinFieldForeign);
-            bob->append("asField", eln->joinField);
+            bob->append("asField", eln->joinField.fullPath());
             bob->append("strategy", EqLookupNode::serializeLookupStrategy(eln->lookupStrategy));
             if (eln->idxEntry) {
                 bob->append("indexName", eln->idxEntry->identifier.catalogName);
diff --git a/src/mongo/db/query/query_planner_test_lib.cpp b/src/mongo/db/query/query_planner_test_lib.cpp
index 8e5b348dcc4..bfb18f47f1e 100644
--- a/src/mongo/db/query/query_planner_test_lib.cpp
+++ b/src/mongo/db/query/query_planner_test_lib.cpp
@@ -1400,12 +1400,12 @@ Status QueryPlannerTestLib::solutionMatches(const BSONObj& testSoln,
                         << testSoln.toString()};
         }
 
-        if (expectedAsField.str() != actualEqLookupNode->joinField) {
+        if (expectedAsField.str() != actualEqLookupNode->joinField.fullPath()) {
             return {ErrorCodes::Error{6267508},
                     str::stream() << "Test solution 'joinField' does not match actual; test "
                                      ""
                                   << expectedAsField.str() << " != actual "
-                                  << actualEqLookupNode->joinField};
+                                  << actualEqLookupNode->joinField.fullPath()};
         }
 
         auto expectedStrategy = expectedEqLookupSoln["strategy"];
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index d3291ba6670..5532fd7cc1c 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -1575,7 +1575,7 @@ void EqLookupNode::appendToString(str::stream* ss, int indent) const {
     addIndent(ss, indent + 1);
     *ss << "from = " << foreignCollection << "\n";
     addIndent(ss, indent + 1);
-    *ss << "as = " << joinField << "\n";
+    *ss << "as = " << joinField.fullPath() << "\n";
     addIndent(ss, indent + 1);
     *ss << "localField = " << joinFieldLocal << "\n";
     addIndent(ss, indent + 1);
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index 04ee6fd9ff1..a22056f09bf 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -1426,7 +1426,7 @@ struct EqLookupNode : public QuerySolutionNode {
                  const std::string& foreignCollection,
                  const std::string& joinFieldLocal,
                  const std::string& joinFieldForeign,
-                 const std::string& joinField)
+                 const FieldPath& joinField)
         : QuerySolutionNode(std::move(child)),
           foreignCollection(foreignCollection),
           joinFieldLocal(joinFieldLocal),
@@ -1485,7 +1485,7 @@ struct EqLookupNode : public QuerySolutionNode {
      * The field stores the array of all matched foreign (inner) documents.
      * If the field already exists in the local (outer) document, the field will be overwritten.
      */
-    std::string joinField;
+    FieldPath joinField;
 
     /**
      * The algorithm that will be used to execute this 'EqLookupNode'. Defaults to nested loop join
diff --git a/src/mongo/db/query/sbe_stage_builder_lookup.cpp b/src/mongo/db/query/sbe_stage_builder_lookup.cpp
index e42c75f868e..1c098e08633 100644
--- a/src/mongo/db/query/sbe_stage_builder_lookup.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_lookup.cpp
@@ -232,6 +232,48 @@ std::pair<SlotId /*matched docs*/, std::unique_ptr<sbe::PlanStage>> buildNljLook
     return {innerResultSlot, std::move(nlj)};
 }
 
+std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildLookupResultObject(
+    std::unique_ptr<sbe::PlanStage> stage,
+    SlotId localDocumentSlot,
+    SlotId resultArraySlot,
+    const FieldPath& fieldPath,
+    const PlanNodeId nodeId,
+    SlotIdGenerator& slotIdGenerator) {
+    const int32_t pathLength = fieldPath.getPathLength();
+
+    // Extract values of all fields along the path except the last one.
+    auto fieldSlots = slotIdGenerator.generateMultiple(pathLength - 1);
+    for (int32_t i = 0; i < pathLength - 1; i++) {
+        const auto fieldName = fieldPath.getFieldName(i);
+        const auto inputSlot = i == 0 ? localDocumentSlot : fieldSlots[i - 1];
+        stage = makeProjectStage(
+            std::move(stage),
+            nodeId,
+            fieldSlots[i],
+            makeFunction("getField"_sd, makeVariable(inputSlot), makeConstant(fieldName)));
+    }
+
+    // Construct new objects for each path level.
+    auto objectSlots = slotIdGenerator.generateMultiple(pathLength);
+    for (int32_t i = pathLength - 1; i >= 0; i--) {
+        const auto rootObjectSlot = i == 0 ? localDocumentSlot : fieldSlots[i - 1];
+        const auto fieldName = fieldPath.getFieldName(i).toString();
+        const auto valueSlot = i == pathLength - 1 ? resultArraySlot : objectSlots[i + 1];
+        stage = makeS<MakeBsonObjStage>(std::move(stage),
+                                        objectSlots[i],                        /* objSlot */
+                                        rootObjectSlot,                        /* rootSlot */
+                                        MakeBsonObjStage::FieldBehavior::drop, /* fieldBehaviour */
+                                        std::vector<std::string>{},            /* fields */
+                                        std::vector<std::string>{fieldName},   /* projectFields */
+                                        SlotVector{valueSlot},                 /* projectVars */
+                                        true,                                  /* forceNewObject */
+                                        false,                                 /* returnOldObject */
+                                        nodeId);
+    }
+
+    return {objectSlots.front(), std::move(stage)};
+}
+
 /*
  * Build $lookup stage using index join strategy. Below is an example plan for the aggregation
  * [{$lookup: {localField: "a", foreignField: "b"}}] with an index {b: 1} on the foreign
@@ -272,7 +314,6 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage(
     std::unique_ptr<sbe::PlanStage> localStage,
     SlotId localRecordSlot,
     std::string localFieldName,
-    std::string joinFieldName,
     const CollectionPtr& foreignColl,
     const IndexEntry& index,
     StringMap<const IndexAccessMethod*>& iamMap,
@@ -394,22 +435,7 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage(
                                          nullptr,
                                          nodeId);
 
-    // TODO(SERVER-63753): Remove mkbson stage here as it's temporarily added to enable testing
-    // index join.
-    auto resultSlot = slotIdGenerator.generate();
-    auto resultStage = makeS<MakeBsonObjStage>(
-        std::move(nljStage),
-        resultSlot,
-        localRecordSlot,
-        MakeObjFieldBehavior::drop /* fieldBehavior */,
-        std::vector<std::string>{} /* fields */,
-        std::vector<std::string>{std::move(joinFieldName)} /* projectFields */,
-        makeSV(foreignGroupSlot) /* projectVars */,
-        true,
-        false,
-        nodeId);
-
-    return {resultSlot, std::move(resultStage)};
+    return {foreignGroupSlot, std::move(nljStage)};
 }
 }  // namespace
 
@@ -419,101 +445,101 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
 
     // $lookup creates its own output documents.
     _shouldProduceRecordIdSlot = false;
-    switch (eqLookupNode->lookupStrategy) {
-        case EqLookupNode::LookupStrategy::kHashJoin:
-            uasserted(5842602, "$lookup planning logic picked hash join");
-            break;
-        case EqLookupNode::LookupStrategy::kIndexedLoopJoin: {
-            tassert(6357201,
+
+    auto localReqs = reqs.copy().set(kResult);
+    auto [localStage, localOutputs] = build(eqLookupNode->children[0], localReqs);
+    SlotId localDocumentSlot = localOutputs.get(PlanStageSlots::kResult);
+
+    auto [matchedDocumentsSlot, foreignStage] = [&, localStage = std::move(localStage)]() mutable
+        -> std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> {
+        switch (eqLookupNode->lookupStrategy) {
+            case EqLookupNode::LookupStrategy::kHashJoin:
+                uasserted(5842602, "$lookup planning logic picked hash join");
+                break;
+            case EqLookupNode::LookupStrategy::kIndexedLoopJoin: {
+                tassert(
+                    6357201,
                     "$lookup using index join should have one child and a populated index entry",
                     eqLookupNode->children.size() == 1 && eqLookupNode->idxEntry);
 
-            const NamespaceString foreignCollNs(eqLookupNode->foreignCollection);
-            const auto& foreignColl = _collections.lookupCollection(foreignCollNs);
-            tassert(6357202,
-                    str::stream() << "$lookup using index join with unknown foreign collection '"
-                                  << foreignCollNs << "'",
-                    foreignColl);
-            const auto& index = *eqLookupNode->idxEntry;
-
-            uassert(6357203,
-                    str::stream() << "$lookup using index join doesn't work for hashed index '"
-                                  << index.identifier.catalogName << "'",
-                    index.type != INDEX_HASHED);
-
-            const auto& localRoot = eqLookupNode->children[0];
-            auto [localStage, localOutputs] = build(localRoot, reqs);
-            sbe::value::SlotId localScanSlot = localOutputs.get(PlanStageSlots::kResult);
-
-            auto [resultSlot, indexJoinStage] =
-                buildIndexJoinLookupStage(_state,
-                                          std::move(localStage),
-                                          localScanSlot,
-                                          eqLookupNode->joinFieldLocal,
-                                          eqLookupNode->joinField,
-                                          foreignColl,
-                                          index,
-                                          _data.iamMap,
-                                          _yieldPolicy,
-                                          eqLookupNode->nodeId(),
-                                          _slotIdGenerator);
-
-            PlanStageSlots outputs;
-            outputs.set(kResult, resultSlot);
-            return {std::move(indexJoinStage), std::move(outputs)};
-        }
-        case EqLookupNode::LookupStrategy::kNestedLoopJoin: {
-            auto numChildren = eqLookupNode->children.size();
-            tassert(6355300, "An EqLookupNode can only have one child", numChildren == 1);
-            const auto& localRoot = eqLookupNode->children[0];
-            auto [localStage, localOutputs] = build(localRoot, reqs);
-            sbe::value::SlotId localResultSlot = localOutputs.get(PlanStageSlots::kResult);
-
-            auto foreignResultSlot = _slotIdGenerator.generate();
-            auto foreignRecordIdSlot = _slotIdGenerator.generate();
-            const auto& foreignColl =
-                _collections.lookupCollection(NamespaceString(eqLookupNode->foreignCollection));
-
-            // TODO SERVER-64091: Delete this tassert when we correctly handle the case of a non
-            //  existent foreign collection.
-            tassert(6355302, "The foreign collection should exist", foreignColl);
-            auto foreignStage = sbe::makeS<sbe::ScanStage>(foreignColl->uuid(),
-                                                           foreignResultSlot,
-                                                           foreignRecordIdSlot,
-                                                           boost::none /* snapshotIdSlot */,
-                                                           boost::none /* indexIdSlot */,
-                                                           boost::none /* indexKeySlot */,
-                                                           boost::none /* indexKeyPatternSlot */,
-                                                           boost::none /* tsSlot */,
-                                                           std::vector<std::string>{} /* fields */,
-                                                           sbe::makeSV() /* vars */,
-                                                           boost::none /* seekKeySlot */,
-                                                           true /* forward */,
-                                                           _yieldPolicy,
-                                                           eqLookupNode->nodeId(),
-                                                           sbe::ScanCallbacks{});
-
-            auto [matchedSlot, nljStage] = buildNljLookupStage(_state,
-                                                               std::move(localStage),
-                                                               localResultSlot,
-                                                               eqLookupNode->joinFieldLocal,
-                                                               std::move(foreignStage),
-                                                               foreignResultSlot,
-                                                               eqLookupNode->joinFieldForeign,
-                                                               eqLookupNode->nodeId(),
-                                                               _slotIdGenerator);
-
-            PlanStageSlots outputs;
-            outputs.set(kResult,
-                        localResultSlot);  // TODO SERVER-63753: create an object for $lookup result
-            outputs.set("local"_sd, localResultSlot);
-            outputs.set("matched"_sd, matchedSlot);
-            return {std::move(nljStage), std::move(outputs)};
+                const NamespaceString foreignCollNs(eqLookupNode->foreignCollection);
+                const auto& foreignColl = _collections.lookupCollection(foreignCollNs);
+                tassert(6357202,
+                        str::stream()
+                            << "$lookup using index join with unknown foreign collection '"
+                            << foreignCollNs << "'",
+                        foreignColl);
+                const auto& index = *eqLookupNode->idxEntry;
+
+                uassert(6357203,
+                        str::stream() << "$lookup using index join doesn't work for hashed index '"
+                                      << index.identifier.catalogName << "'",
+                        index.type != INDEX_HASHED);
+
+                return buildIndexJoinLookupStage(_state,
+                                                 std::move(localStage),
+                                                 localDocumentSlot,
+                                                 eqLookupNode->joinFieldLocal,
+                                                 foreignColl,
+                                                 index,
+                                                 _data.iamMap,
+                                                 _yieldPolicy,
+                                                 eqLookupNode->nodeId(),
+                                                 _slotIdGenerator);
+            }
+            case EqLookupNode::LookupStrategy::kNestedLoopJoin: {
+                auto numChildren = eqLookupNode->children.size();
+                tassert(6355300, "An EqLookupNode can only have one child", numChildren == 1);
+
+                auto foreignResultSlot = _slotIdGenerator.generate();
+                auto foreignRecordIdSlot = _slotIdGenerator.generate();
+                const auto& foreignColl =
+                    _collections.lookupCollection(NamespaceString(eqLookupNode->foreignCollection));
+
+                // TODO SERVER-64091: Delete this tassert when we correctly handle the case of a non
+                //  existent foreign collection.
+                tassert(6355302, "The foreign collection should exist", foreignColl);
+                auto foreignStage = makeS<ScanStage>(foreignColl->uuid(),
+                                                     foreignResultSlot,
+                                                     foreignRecordIdSlot,
+                                                     boost::none /* snapshotIdSlot */,
+                                                     boost::none /* indexIdSlot */,
+                                                     boost::none /* indexKeySlot */,
+                                                     boost::none /* indexKeyPatternSlot */,
+                                                     boost::none /* tsSlot */,
+                                                     std::vector<std::string>{} /* fields */,
+                                                     makeSV() /* vars */,
+                                                     boost::none /* seekKeySlot */,
+                                                     true /* forward */,
+                                                     _yieldPolicy,
+                                                     eqLookupNode->nodeId(),
+                                                     ScanCallbacks{});
+
+                return buildNljLookupStage(_state,
+                                           std::move(localStage),
+                                           localDocumentSlot,
+                                           eqLookupNode->joinFieldLocal,
+                                           std::move(foreignStage),
+                                           foreignResultSlot,
+                                           eqLookupNode->joinFieldForeign,
+                                           eqLookupNode->nodeId(),
+                                           _slotIdGenerator);
+            }
+            default:
+                MONGO_UNREACHABLE_TASSERT(5842605);
         }
-        default:
-            MONGO_UNREACHABLE_TASSERT(5842605);
-    }
-    MONGO_UNREACHABLE_TASSERT(5842606);
+    }();
+
+    auto [resultSlot, resultStage] = buildLookupResultObject(std::move(foreignStage),
+                                                             localDocumentSlot,
+                                                             matchedDocumentsSlot,
+                                                             eqLookupNode->joinField,
+                                                             eqLookupNode->nodeId(),
+                                                             _slotIdGenerator);
+
+    PlanStageSlots outputs;
+    outputs.set(kResult, resultSlot);
+    return {std::move(resultStage), std::move(outputs)};
 }
 
 }  // namespace mongo::stage_builder
diff --git a/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp b/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp
index 055b685e684..8cb7d76fc2d 100644
--- a/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp
@@ -64,185 +64,166 @@ public:
         // next opTime (LocalOplogInfo::getNextOpTimes) to use for a write.
         repl::createOplog(opCtx());
 
-        // Acquire the lock for our inner collection in MODE_X as we will perform writes.
-        uassertStatusOK(_storage->createCollection(opCtx(), _secondaryNss, CollectionOptions()));
+        // Create local and foreign collections.
+        ASSERT_OK(_storage->createCollection(opCtx(), _nss, CollectionOptions()));
+        ASSERT_OK(_storage->createCollection(opCtx(), _foreignNss, CollectionOptions()));
     }
 
     virtual void tearDown() {
         _storage.reset();
-        _secondaryCollLock.reset();
+        _localCollLock.reset();
+        _foreignCollLock.reset();
         SbeStageBuilderTestFixture::tearDown();
     }
 
-    // VirtualScanNode wants the docs wrapped into BSONArray.
-    std::vector<BSONArray> prepInputForVirtualScanNode(const std::vector<BSONObj>& docs) {
-        std::vector<BSONArray> input;
-        input.reserve(docs.size());
-        for (const auto& doc : docs) {
-            input.emplace_back(BSON_ARRAY(doc));
-        }
-        return input;
-    }
-
-    // Constructs a QuerySolution consisting of a EqLookupNode on top of two VirtualScanNodes.
-    std::unique_ptr<QuerySolution> makeLookupSolution(const std::string& lookupSpec,
-                                                      const std::string& fromColl,
-                                                      const std::vector<BSONObj>& localDocs,
-                                                      const std::vector<BSONObj>& foreignDocs) {
-        auto expCtx = make_intrusive<ExpressionContextForTest>();
-        auto lookupNss = NamespaceString{"test", fromColl};
-        expCtx->setResolvedNamespace(lookupNss,
-                                     ExpressionContext::ResolvedNamespace{lookupNss, {}});
-
-        auto docSource =
-            DocumentSourceLookUp::createFromBson(fromjson(lookupSpec).firstElement(), expCtx);
-        auto docLookup = static_cast<DocumentSourceLookUp*>(docSource.get());
-
-        auto localVirtScanNode =
-            std::make_unique<VirtualScanNode>(prepInputForVirtualScanNode(localDocs),
-                                              VirtualScanNode::ScanType::kCollScan,
-                                              false /*hasRecordId*/);
-
-        auto lookupNode = std::make_unique<EqLookupNode>(std::move(localVirtScanNode),
-                                                         fromColl,
-                                                         docLookup->getLocalField()->fullPath(),
-                                                         docLookup->getForeignField()->fullPath(),
-                                                         docLookup->getAsField().fullPath());
-
-        std::vector<InsertStatement> inserts;
-        for (const auto& doc : foreignDocs) {
-            inserts.emplace_back(doc);
-        }
-
-        // Perform our writes by acquiring the lock in MODE_X.
-        _secondaryCollLock =
-            std::make_unique<AutoGetCollection>(opCtx(), _secondaryNss, LockMode::MODE_X);
+    void insertDocuments(const NamespaceString& nss,
+                         std::unique_ptr<AutoGetCollection>& lock,
+                         const std::vector<BSONObj>& docs) {
+        std::vector<InsertStatement> inserts{docs.begin(), docs.end()};
+        lock = std::make_unique<AutoGetCollection>(opCtx(), nss, LockMode::MODE_X);
         {
             WriteUnitOfWork wuow{opCtx()};
-            uassertStatusOK(
-                _secondaryCollLock.get()->getWritableCollection(opCtx())->insertDocuments(
-                    opCtx(), inserts.begin(), inserts.end(), nullptr /* opDebug */));
+            ASSERT_OK(lock.get()->getWritableCollection(opCtx())->insertDocuments(
+                opCtx(), inserts.begin(), inserts.end(), nullptr /* opDebug */));
             wuow.commit();
         }
-        _secondaryCollLock.reset();
 
         // Before we read, lock the collection in MODE_IS.
-        _secondaryCollLock =
-            std::make_unique<AutoGetCollection>(opCtx(), _secondaryNss, LockMode::MODE_IS);
+        lock = std::make_unique<AutoGetCollection>(opCtx(), nss, LockMode::MODE_IS);
+    }
+
+    void insertDocuments(const std::vector<BSONObj>& localDocs,
+                         const std::vector<BSONObj>& foreignDocs) {
+        insertDocuments(_nss, _localCollLock, localDocs);
+        insertDocuments(_foreignNss, _foreignCollLock, foreignDocs);
 
-        // While the main collection does not exist because the input to an EqLookupNode in these
-        // tests is a VirtualScanNode, we need a real collection for the foreign side.
         _collections = MultipleCollectionAccessor(opCtx(),
-                                                  nullptr /* mainColl */,
+                                                  &_localCollLock->getCollection(),
                                                   _nss,
                                                   false /* isAnySecondaryNamespaceAViewOrSharded */,
-                                                  {_secondaryNss});
-        return makeQuerySolution(std::move(lookupNode));
+                                                  {_foreignNss});
     }
 
-    // Execute the stage tree and check the results.
-    void CheckNljResults(PlanStage* nljStage,
-                         SlotId localSlot,
-                         SlotId matchedSlot,
-                         const std::vector<std::pair<BSONObj, std::vector<BSONObj>>>& expected,
-                         bool debugPrint = false) {
+    struct CompiledTree {
+        std::unique_ptr<sbe::PlanStage> stage;
+        stage_builder::PlanStageData data;
+        std::unique_ptr<CompileCtx> ctx;
+        SlotAccessor* resultSlotAccessor;
+    };
+
+    // Constructs ready-to-execute SBE tree for $lookup specified by the arguments.
+    CompiledTree buildLookupSbeTree(const std::string& localKey,
+                                    const std::string& foreignKey,
+                                    const std::string& asKey) {
+        // Documents from the local collection are provided using collection scan.
+        auto localScanNode = std::make_unique<CollectionScanNode>();
+        localScanNode->name = _nss.toString();
+
+        // Construct logical query solution.
+        auto foreignCollName = _foreignNss.toString();
+        auto lookupNode = std::make_unique<EqLookupNode>(
+            std::move(localScanNode), foreignCollName, localKey, foreignKey, asKey);
+        auto solution = makeQuerySolution(std::move(lookupNode));
+
+        // Convert logical solution into the physical SBE plan.
+        auto [resultSlots, stage, data, _] = buildPlanStage(std::move(solution),
+                                                            false /*hasRecordId*/,
+                                                            nullptr /*shard filterer*/,
+                                                            nullptr /*collator*/);
 
+        // Prepare the SBE tree for execution.
         auto ctx = makeCompileCtx();
-        prepareTree(ctx.get(), nljStage);
-        SlotAccessor* outer = nljStage->getAccessor(*ctx, localSlot);
-        SlotAccessor* inner = nljStage->getAccessor(*ctx, matchedSlot);
+        prepareTree(ctx.get(), stage.get());
+
+        auto resultSlot = data.outputs.get(stage_builder::PlanStageSlots::kResult);
+        SlotAccessor* resultSlotAccessor = stage->getAccessor(*ctx, resultSlot);
+
+        return CompiledTree{.stage = std::move(stage),
+                            .data = std::move(data),
+                            .ctx = std::move(ctx),
+                            .resultSlotAccessor = resultSlotAccessor};
+    }
+
+    // Check that SBE plan for '$lookup' returns expected documents.
+    void assertReturnedDocuments(const std::string& localKey,
+                                 const std::string& foreignKey,
+                                 const std::string& asKey,
+                                 const std::vector<BSONObj>& expected,
+                                 bool debugPrint = false) {
+        auto tree = buildLookupSbeTree(localKey, foreignKey, asKey);
+        auto& stage = tree.stage;
+
+        if (debugPrint) {
+            std::cout << std::endl << DebugPrinter{true}.print(stage->debugPrint()) << std::endl;
+        }
 
         size_t i = 0;
-        for (auto st = nljStage->getNext(); st == PlanState::ADVANCED;
-             st = nljStage->getNext(), i++) {
-            auto [outerTag, outerVal] = outer->copyOrMoveValue();
-            ValueGuard outerGuard{outerTag, outerVal};
+        for (auto state = stage->getNext(); state == PlanState::ADVANCED;
+             state = stage->getNext(), i++) {
+            // Retrieve the result document from SBE plan.
+            auto [resultTag, resultValue] = tree.resultSlotAccessor->copyOrMoveValue();
+            ValueGuard resultGuard{resultTag, resultValue};
             if (debugPrint) {
-                std::cout << i << " outer: " << std::make_pair(outerTag, outerVal) << std::endl;
+                std::cout << "Actual document: " << std::make_pair(resultTag, resultValue)
+                          << std::endl;
             }
+
+            // If the plan returned more documents than expected, proceed extracting all of them.
+            // This way, the developer will see them if debug print is enabled.
             if (i >= expected.size()) {
-                // We'll assert eventually that there were more actual results than expected.
                 continue;
             }
 
-            auto [expectedOuterTag, expectedOuterVal] = copyValue(
-                TypeTags::bsonObject, bitcastFrom<const char*>(expected[i].first.objdata()));
-            ValueGuard expectedOuterGuard{expectedOuterTag, expectedOuterVal};
-
-            assertValuesEqual(outerTag, outerVal, expectedOuterTag, expectedOuterVal);
-
-            auto [innerTag, innerVal] = inner->copyOrMoveValue();
-            ValueGuard innerGuard{innerTag, innerVal};
-
-            ASSERT_EQ(innerTag, TypeTags::Array);
-            auto innerMatches = getArrayView(innerVal);
-
+            // Construct view to the expected document.
+            auto [expectedTag, expectedValue] =
+                copyValue(TypeTags::bsonObject, bitcastFrom<const char*>(expected[i].objdata()));
             if (debugPrint) {
-                std::cout << "  inner:" << std::endl;
-                for (size_t m = 0; m < innerMatches->size(); m++) {
-                    auto [matchedTag, matchedVal] = innerMatches->getAt(m);
-                    std::cout << "  " << m << ": " << std::make_pair(matchedTag, matchedVal)
-                              << std::endl;
-                }
+                std::cout << "Expected document: " << std::make_pair(expectedTag, expectedValue)
+                          << std::endl;
             }
 
-            ASSERT_EQ(innerMatches->size(), expected[i].second.size());
-            for (size_t m = 0; m < innerMatches->size(); m++) {
-                auto [matchedTag, matchedVal] = innerMatches->getAt(m);
-                auto [expectedMatchTag, expectedMatchVal] =
-                    copyValue(TypeTags::bsonObject,
-                              bitcastFrom<const char*>(expected[i].second[m].objdata()));
-                ValueGuard expectedMatchGuard{expectedMatchTag, expectedMatchVal};
-                assertValuesEqual(matchedTag, matchedVal, expectedMatchTag, expectedMatchVal);
-            }
+            // Assert that the document from SBE plan is equal to the expected one.
+            assertValuesEqual(resultTag, resultValue, expectedTag, expectedValue);
         }
+
         ASSERT_EQ(i, expected.size());
-        nljStage->close();
+        stage->close();
     }
 
-    void runTest(const std::vector<BSONObj>& ldocs,
-                 const std::vector<BSONObj>& fdocs,
-                 const std::string& lkey,
-                 const std::string& fkey,
-                 const std::vector<std::pair<BSONObj, std::vector<BSONObj>>>& expected,
-                 bool debugPrint = false) {
-        const char* foreignCollName = _secondaryNss.toString().data();
-        std::stringstream lookupSpec;
-        lookupSpec << "{$lookup: ";
-        lookupSpec << "  {";
-        lookupSpec << "    from: '" << foreignCollName << "', ";
-        lookupSpec << "    localField:   '" << lkey << "', ";
-        lookupSpec << "    foreignField: '" << fkey << "', ";
-        lookupSpec << "    as: 'matched'";
-        lookupSpec << "  }";
-        lookupSpec << "}";
-
-        auto solution = makeLookupSolution(lookupSpec.str(), foreignCollName, ldocs, fdocs);
-        auto [resultSlots, stage, data, _] = buildPlanStage(std::move(solution),
-                                                            false /*hasRecordId*/,
-                                                            nullptr /*shard filterer*/,
-                                                            nullptr /*collator*/);
-        if (debugPrint) {
-            debugPrintPlan(*stage);
+    // Check that SBE plan for '$lookup' returns expected documents. Expected documents are
+    // described in pairs '(local document, matched foreign documents)'.
+    void assertMatchedDocuments(
+        const std::string& localKey,
+        const std::string& foreignKey,
+        const std::vector<std::pair<BSONObj, std::vector<BSONObj>>>& expectedPairs,
+        bool debugPrint = false) {
+        const std::string resultFieldName{"result"};
+
+        // Construct expected documents.
+        std::vector<BSONObj> expectedDocuments;
+        expectedDocuments.reserve(expectedPairs.size());
+        for (auto& [localDocument, matchedDocuments] : expectedPairs) {
+            MutableDocument expectedDocument;
+            expectedDocument.reset(localDocument, false /* stripMetadata */);
+
+            std::vector<mongo::Value> matchedValues{matchedDocuments.begin(),
+                                                    matchedDocuments.end()};
+            expectedDocument.setField(resultFieldName, mongo::Value{matchedValues});
+            const auto expectedBson = expectedDocument.freeze().toBson();
+
+            expectedDocuments.push_back(expectedBson);
         }
 
-        CheckNljResults(stage.get(),
-                        data.outputs.get("local"),
-                        data.outputs.get("matched"),
-                        expected,
-                        debugPrint);
-    }
-
-    void debugPrintPlan(const PlanStage& stage, StringData header = "") {
-        std::cout << std::endl << "*** " << header << " ***" << std::endl;
-        std::cout << DebugPrinter{}.print(stage.debugPrint());
-        std::cout << std::endl;
+        assertReturnedDocuments(
+            localKey, foreignKey, resultFieldName, expectedDocuments, debugPrint);
     }
 
 private:
-    const NamespaceString _secondaryNss = NamespaceString{"testdb.sbe_stage_builder_secondary"};
     std::unique_ptr<repl::StorageInterface> _storage;
-    std::unique_ptr<AutoGetCollection> _secondaryCollLock = nullptr;
+
+    const NamespaceString _foreignNss{"testdb.sbe_stage_builder_foreign"};
+    std::unique_ptr<AutoGetCollection> _localCollLock = nullptr;
+    std::unique_ptr<AutoGetCollection> _foreignCollLock = nullptr;
 };
 
 TEST_F(LookupStageBuilderTest, NestedLoopJoin_Basic) {
@@ -273,7 +254,8 @@ TEST_F(LookupStageBuilderTest, NestedLoopJoin_Basic) {
         {ldocs[3], {fdocs[0], fdocs[2], fdocs[3]}},
     };
 
-    runTest(ldocs, fdocs, "lkey", "fkey", expected);
+    insertDocuments(ldocs, fdocs);
+    assertMatchedDocuments("lkey", "fkey", expected);
 }
 
 TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_Null) {
@@ -292,7 +274,8 @@ TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_Null) {
         {ldocs[0], {fdocs[1], fdocs[2], fdocs[3], fdocs[4], fdocs[5]}},
     };
 
-    runTest(ldocs, fdocs, "lkey", "fkey", expected);
+    insertDocuments(ldocs, fdocs);
+    assertMatchedDocuments("lkey", "fkey", expected);
 }
 
 TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_Missing) {
@@ -311,7 +294,8 @@ TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_Missing) {
         {ldocs[0], {fdocs[1], fdocs[2], fdocs[3], fdocs[4], fdocs[5]}},
     };
 
-    runTest(ldocs, fdocs, "lkey", "fkey", expected);
+    insertDocuments(ldocs, fdocs);
+    assertMatchedDocuments("lkey", "fkey", expected);
 }
 
 TEST_F(LookupStageBuilderTest, NestedLoopJoin_EmptyArrays) {
@@ -335,7 +319,8 @@ TEST_F(LookupStageBuilderTest, NestedLoopJoin_EmptyArrays) {
         {ldocs[1], {fdocs[7]}},  // TODO SEVER-63700: it should be {fdocs[6], fdocs[7]}
     };
 
-    runTest(ldocs, fdocs, "lkey", "fkey", expected);
+    insertDocuments(ldocs, fdocs);
+    assertMatchedDocuments("lkey", "fkey", expected);
 }
 
 TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_SubFieldScalar) {
@@ -362,7 +347,8 @@ TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_SubFieldScalar) {
     };
 
     // TODO SERVER-63690: enable this test.
-    // runTest(ldocs, fdocs, "nested.lkey", "fkey", expected);
+    // insertDocuments(ldocs, fdocs);
+    // assertMatchedDocuments("nested.lkey", "fkey", expected);
 }
 
 TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_SubFieldArray) {
@@ -399,7 +385,8 @@ TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_SubFieldArray) {
     };
 
     // TODO SERVER-63690: enable this test.
-    // runTest(ldocs, fdocs, "nested.lkey", "fkey", expected, true);
+    // insertDocuments(ldocs, fdocs);
+    // assertMatchedDocuments("nested.lkey", "fkey", expected, true);
 }
 
 TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_PathWithNumber) {
@@ -429,6 +416,75 @@ TEST_F(LookupStageBuilderTest, NestedLoopJoin_LocalKey_PathWithNumber) {
     };
 
     // TODO SERVER-63690: either remove or enable this test.
-    // runTest(ldocs, fdocs, "nested.0.lkey", "fkey", expected, true);
+    // insertDocuments(ldocs, fdocs);
+    // assertMatchedDocuments("nested.0.lkey", "fkey", expected, true);
+}
+
+TEST_F(LookupStageBuilderTest, OneComponentAsPath) {
+    insertDocuments({fromjson("{_id: 0}")}, {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments("_id", "_id", "result", {fromjson("{_id: 0, result: [{_id: 0}]}")});
+}
+
+TEST_F(LookupStageBuilderTest, OneComponentAsPathReplacingExistingObject) {
+    insertDocuments({fromjson("{_id: 0, result: {a: {b: 1}, c: 2}}")}, {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments("_id", "_id", "result", {fromjson("{_id: 0, result: [{_id: 0}]}")});
+}
+
+TEST_F(LookupStageBuilderTest, OneComponentAsPathReplacingExistingArray) {
+    insertDocuments({fromjson("{_id: 0, result: [{a: 1}, {b: 2}]}")}, {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments("_id", "_id", "result", {fromjson("{_id: 0, result: [{_id: 0}]}")});
+}
+
+TEST_F(LookupStageBuilderTest, ThreeComponentAsPath) {
+    insertDocuments({fromjson("{_id: 0}")}, {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments(
+        "_id", "_id", "one.two.three", {fromjson("{_id: 0, one: {two: {three: [{_id: 0}]}}}")});
+}
+
+TEST_F(LookupStageBuilderTest, ThreeComponentAsPathExtendingExistingObjectOnOneLevel) {
+    insertDocuments({fromjson("{_id: 0, one: {a: 1}}")}, {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments("_id",
+                            "_id",
+                            "one.two.three",
+                            {fromjson("{_id: 0, one: {a: 1, two: {three: [{_id: 0}]}}}")});
+}
+
+TEST_F(LookupStageBuilderTest, ThreeComponentAsPathExtendingExistingObjectOnTwoLevels) {
+    insertDocuments({fromjson("{_id: 0, one: {a: 1, two: {b: 2}}}")}, {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments("_id",
+                            "_id",
+                            "one.two.three",
+                            {fromjson("{_id: 0, one: {a: 1, two: {b: 2, three: [{_id: 0}]}}}")});
+}
+
+TEST_F(LookupStageBuilderTest, ThreeComponentAsPathReplacingSingleValueInExistingObject) {
+    insertDocuments({fromjson("{_id: 0, one: {a: 1, two: {b: 2, three: 3}}}}")},
+                    {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments("_id",
+                            "_id",
+                            "one.two.three",
+                            {fromjson("{_id: 0, one: {a: 1, two: {b: 2, three: [{_id: 0}]}}}")});
+}
+
+TEST_F(LookupStageBuilderTest, ThreeComponentAsPathReplacingExistingArray) {
+    insertDocuments({fromjson("{_id: 0, one: [{a: 1}, {b: 2}]}")}, {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments(
+        "_id", "_id", "one.two.three", {fromjson("{_id: 0, one: {two: {three: [{_id: 0}]}}}")});
+}
+
+TEST_F(LookupStageBuilderTest, ThreeComponentAsPathDoesNotPerformArrayTraversal) {
+    insertDocuments({fromjson("{_id: 0, one: [{a: 1, two: [{b: 2, three: 3}]}]}")},
+                    {fromjson("{_id: 0}")});
+
+    assertReturnedDocuments(
+        "_id", "_id", "one.two.three", {fromjson("{_id: 0, one: {two: {three: [{_id: 0}]}}}")});
 }
 }  // namespace mongo::sbe