From e09abbd7289641317d9e213204fb79731655e004 Mon Sep 17 00:00:00 2001 From: Nikita Lapkov Date: Tue, 15 Mar 2022 11:28:17 +0000 Subject: SERVER-63753 Translate $lookup result object creation in SBE --- src/mongo/db/query/plan_explainer_sbe.cpp | 2 +- src/mongo/db/query/query_planner_test_lib.cpp | 4 +- src/mongo/db/query/query_solution.cpp | 2 +- src/mongo/db/query/query_solution.h | 4 +- src/mongo/db/query/sbe_stage_builder_lookup.cpp | 242 ++++++++------- .../db/query/sbe_stage_builder_lookup_test.cpp | 344 ++++++++++++--------- 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> buildNljLook return {innerResultSlot, std::move(nlj)}; } +std::pair> buildLookupResultObject( + std::unique_ptr 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(std::move(stage), + objectSlots[i], /* objSlot */ + rootObjectSlot, /* rootSlot */ + MakeBsonObjStage::FieldBehavior::drop, /* fieldBehaviour */ + std::vector{}, /* fields */ + std::vector{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> buildIndexJoinLookupStage( std::unique_ptr localStage, SlotId localRecordSlot, std::string localFieldName, - std::string joinFieldName, const CollectionPtr& foreignColl, const IndexEntry& index, StringMap& iamMap, @@ -394,22 +435,7 @@ std::pair> 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( - std::move(nljStage), - resultSlot, - localRecordSlot, - MakeObjFieldBehavior::drop /* fieldBehavior */, - std::vector{} /* fields */, - std::vector{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, 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> { + 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(foreignColl->uuid(), - foreignResultSlot, - foreignRecordIdSlot, - boost::none /* snapshotIdSlot */, - boost::none /* indexIdSlot */, - boost::none /* indexKeySlot */, - boost::none /* indexKeyPatternSlot */, - boost::none /* tsSlot */, - std::vector{} /* 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(foreignColl->uuid(), + foreignResultSlot, + foreignRecordIdSlot, + boost::none /* snapshotIdSlot */, + boost::none /* indexIdSlot */, + boost::none /* indexKeySlot */, + boost::none /* indexKeyPatternSlot */, + boost::none /* tsSlot */, + std::vector{} /* 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 prepInputForVirtualScanNode(const std::vector& docs) { - std::vector 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 makeLookupSolution(const std::string& lookupSpec, - const std::string& fromColl, - const std::vector& localDocs, - const std::vector& foreignDocs) { - auto expCtx = make_intrusive(); - auto lookupNss = NamespaceString{"test", fromColl}; - expCtx->setResolvedNamespace(lookupNss, - ExpressionContext::ResolvedNamespace{lookupNss, {}}); - - auto docSource = - DocumentSourceLookUp::createFromBson(fromjson(lookupSpec).firstElement(), expCtx); - auto docLookup = static_cast(docSource.get()); - - auto localVirtScanNode = - std::make_unique(prepInputForVirtualScanNode(localDocs), - VirtualScanNode::ScanType::kCollScan, - false /*hasRecordId*/); - - auto lookupNode = std::make_unique(std::move(localVirtScanNode), - fromColl, - docLookup->getLocalField()->fullPath(), - docLookup->getForeignField()->fullPath(), - docLookup->getAsField().fullPath()); - - std::vector inserts; - for (const auto& doc : foreignDocs) { - inserts.emplace_back(doc); - } - - // Perform our writes by acquiring the lock in MODE_X. - _secondaryCollLock = - std::make_unique(opCtx(), _secondaryNss, LockMode::MODE_X); + void insertDocuments(const NamespaceString& nss, + std::unique_ptr& lock, + const std::vector& docs) { + std::vector inserts{docs.begin(), docs.end()}; + lock = std::make_unique(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(opCtx(), _secondaryNss, LockMode::MODE_IS); + lock = std::make_unique(opCtx(), nss, LockMode::MODE_IS); + } + + void insertDocuments(const std::vector& localDocs, + const std::vector& 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>>& expected, - bool debugPrint = false) { + struct CompiledTree { + std::unique_ptr stage; + stage_builder::PlanStageData data; + std::unique_ptr 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(); + localScanNode->name = _nss.toString(); + + // Construct logical query solution. + auto foreignCollName = _foreignNss.toString(); + auto lookupNode = std::make_unique( + 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& 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(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(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(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& ldocs, - const std::vector& fdocs, - const std::string& lkey, - const std::string& fkey, - const std::vector>>& 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>>& expectedPairs, + bool debugPrint = false) { + const std::string resultFieldName{"result"}; + + // Construct expected documents. + std::vector expectedDocuments; + expectedDocuments.reserve(expectedPairs.size()); + for (auto& [localDocument, matchedDocuments] : expectedPairs) { + MutableDocument expectedDocument; + expectedDocument.reset(localDocument, false /* stripMetadata */); + + std::vector 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 _storage; - std::unique_ptr _secondaryCollLock = nullptr; + + const NamespaceString _foreignNss{"testdb.sbe_stage_builder_foreign"}; + std::unique_ptr _localCollLock = nullptr; + std::unique_ptr _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 -- cgit v1.2.1