diff options
-rw-r--r-- | jstests/aggregation/sources/lookup/lookup_query_stats.js | 14 | ||||
-rw-r--r-- | jstests/noPassthrough/lookup_max_intermediate_size.js | 10 | ||||
-rw-r--r-- | jstests/noPassthrough/lookup_pushdown.js | 202 | ||||
-rw-r--r-- | jstests/noPassthrough/plan_cache_replan_group_lookup.js | 20 | ||||
-rw-r--r-- | src/mongo/db/pipeline/pipeline_d.cpp | 31 | ||||
-rw-r--r-- | src/mongo/db/query/get_executor.cpp | 3 | ||||
-rw-r--r-- | src/mongo/db/query/planner_analysis.cpp | 56 | ||||
-rw-r--r-- | src/mongo/db/query/planner_analysis.h | 14 | ||||
-rw-r--r-- | src/mongo/db/query/query_planner.cpp | 16 | ||||
-rw-r--r-- | src/mongo/db/query/query_solution.cpp | 2 | ||||
-rw-r--r-- | src/mongo/db/query/query_solution.h | 4 | ||||
-rw-r--r-- | src/mongo/db/query/query_solution_test.cpp | 86 | ||||
-rw-r--r-- | src/mongo/db/query/sbe_runtime_planner.h | 6 | ||||
-rw-r--r-- | src/mongo/db/query/sbe_stage_builder_lookup_test.cpp | 11 |
14 files changed, 392 insertions, 83 deletions
diff --git a/jstests/aggregation/sources/lookup/lookup_query_stats.js b/jstests/aggregation/sources/lookup/lookup_query_stats.js index 46a337ddff7..77731d1df84 100644 --- a/jstests/aggregation/sources/lookup/lookup_query_stats.js +++ b/jstests/aggregation/sources/lookup/lookup_query_stats.js @@ -22,7 +22,8 @@ load("jstests/libs/sbe_explain_helpers.js"); // For getSbePlanStages and // getQueryInfoAtTopLevelOrFirstStage. const isSBELookupEnabled = checkSBEEnabled(db, ["featureFlagSBELookupPushdown"]); - +const isSBELookupNLJEnabled = + checkSBEEnabled(db, ["featureFlagSBELookupPushdown", "featureFlagSbeFull"]); const testDB = db.getSiblingDB("lookup_query_stats"); testDB.dropDatabase(); @@ -87,7 +88,14 @@ let getCurrentQueryExecutorStats = function() { let checkExplainOutputForVerLevel = function( explainOutput, expected, verbosityLevel, expectedQueryPlan) { const lkpStages = getAggPlanStages(explainOutput, "$lookup"); - if (isSBELookupEnabled) { + + // Only make SBE specific assertions when we know that our $lookup has been pushed down. + // More precisely, $lookup pushdown must be enabled, and we must have NLJ pushed down + // enabled or be targeting a different $lookup strategy. + if (isSBELookupEnabled && + (isSBELookupNLJEnabled || + (expectedQueryPlan.hasOwnProperty("strategy") && + expectedQueryPlan.strategy !== "NestedLoopJoin"))) { // If the SBE lookup is enabled, the $lookup stage is pushed down to the SBE and it's // not visible in 'stages' field of the explain output. Instead, 'queryPlan.stage' must be // "EQ_LOOKUP". @@ -202,7 +210,7 @@ let testQueryExecutorStatsWithCollectionScan = function() { // There is no index in the collection. assert.eq(0, curScannedKeys); - if (isSBELookupEnabled) { + if (isSBELookupNLJEnabled) { checkExplainOutputForAllVerbosityLevels( localColl, fromColl, diff --git a/jstests/noPassthrough/lookup_max_intermediate_size.js b/jstests/noPassthrough/lookup_max_intermediate_size.js index b9534b29d38..f2d5704c7da 100644 --- a/jstests/noPassthrough/lookup_max_intermediate_size.js +++ b/jstests/noPassthrough/lookup_max_intermediate_size.js @@ -87,15 +87,15 @@ function runTest(coll, from, expectedErrorCode) { /** * Run tests on single node. */ -const standalone = MongoRunner.runMongod(); +const standalone = MongoRunner.runMongod( + {setParameter: {internalLookupStageIntermediateDocumentMaxSizeBytes: 30 * 1024 * 1024}}); const db = standalone.getDB("test"); -assert.commandWorked(db.adminCommand( - {setParameter: 1, internalLookupStageIntermediateDocumentMaxSizeBytes: 30 * 1024 * 1024})); - db.lookUp.drop(); const expectedErrorCode = - (checkSBEEnabled(db, ["featureFlagSBELookupPushdown"])) ? ErrorCodes.ExceededMemoryLimit : 4568; + (checkSBEEnabled(db, ["featureFlagSBELookupPushdown", "featureFlagSbeFull"])) + ? ErrorCodes.ExceededMemoryLimit + : 4568; runTest(db.lookUp, db.from, expectedErrorCode); MongoRunner.stopMongod(standalone); diff --git a/jstests/noPassthrough/lookup_pushdown.js b/jstests/noPassthrough/lookup_pushdown.js index 25b5d65dbfb..9497e9d135c 100644 --- a/jstests/noPassthrough/lookup_pushdown.js +++ b/jstests/noPassthrough/lookup_pushdown.js @@ -7,7 +7,7 @@ "use strict"; load("jstests/libs/sbe_util.js"); // For 'checkSBEEnabled()'. -load("jstests/libs/analyze_plan.js"); // For 'getAggPlanStages()' and 'hasRejectedPlans()' +load("jstests/libs/analyze_plan.js"); // For 'getAggPlanStages' and other explain helpers. const JoinAlgorithm = { Classic: 0, @@ -18,8 +18,10 @@ const JoinAlgorithm = { }; // Standalone cases. -const conn = MongoRunner.runMongod( - {setParameter: {featureFlagSBELookupPushdown: true, allowDiskUseByDefault: false}}); +const conn = MongoRunner.runMongod({ + setParameter: + {featureFlagSBELookupPushdown: true, featureFlagSbeFull: true, allowDiskUseByDefault: false} +}); assert.neq(null, conn, "mongod was unable to start up"); const name = "lookup_pushdown"; const foreignCollName = "foreign_lookup_pushdown"; @@ -307,9 +309,9 @@ function setLookupPushdownDisabled(value) { const explain = coll.explain().aggregate(pipeline); assert(explain.hasOwnProperty("explainVersion"), explain); if (isSBE) { - assert.eq(explain.explainVersion, 2, explain); + assert.eq(explain.explainVersion, "2", explain); } else { - assert.eq(explain.explainVersion, 1, explain); + assert.eq(explain.explainVersion, "1", explain); } }; @@ -355,6 +357,13 @@ function setLookupPushdownDisabled(value) { runTest(coll, [{$lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "out"}}], JoinAlgorithm.NLJ /* expectedJoinAlgorithm */); + + // If we add an index that is not a partial index, we should then use INLJ. + assert.commandWorked(foreignColl.createIndex({b: 1, a: 1})); + runTest(coll, + [{$lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "out"}}], + JoinAlgorithm.INLJ /* expectedJoinAlgorithm */, + {b: 1, a: 1} /* indexKeyPattern */); assert.commandWorked(foreignColl.dropIndexes()); })(); @@ -378,6 +387,54 @@ function setLookupPushdownDisabled(value) { runTest(coll, [{$lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "out"}}], JoinAlgorithm.NLJ /* expectedJoinAlgorithm */); + + // Insert a document with multikey paths in the foreign collection that will be used for testing + // wildcard indexes. + const mkDoc = {b: [3, 4], c: [5, 6, {d: [7, 8]}]}; + assert.commandWorked(foreignColl.insert(mkDoc)); + + // An incompatible wildcard index should result in using NLJ. + assert.commandWorked(foreignColl.dropIndexes()); + assert.commandWorked(foreignColl.createIndex({'b.$**': 1})); + runTest(coll, + [{ + $lookup: + {from: foreignCollName, localField: "a", foreignField: "not a match", as: "out"} + }], + JoinAlgorithm.NLJ /* expectedJoinAlgorithm */); + + // A compatible wildcard index with no other SBE compatible indexes should result in NLJ. + runTest(coll, + [{$lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "out"}}], + JoinAlgorithm.NLJ /* expectedJoinAlgorithm */); + + runTest(coll, + [{$lookup: {from: foreignCollName, localField: "a", foreignField: "b.c", as: "out"}}], + JoinAlgorithm.NLJ /* expectedJoinAlgorithm */); + + assert.commandWorked(foreignColl.dropIndexes()); + assert.commandWorked(foreignColl.createIndex({'$**': 1}, {wildcardProjection: {b: 1}})); + runTest(coll, + [{$lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "out"}}], + JoinAlgorithm.NLJ /* expectedJoinAlgorithm */); + + // Create a regular index over the foreignField. We should now use INLJ. + assert.commandWorked(foreignColl.createIndex({b: 1})); + runTest(coll, + [{$lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "out"}}], + JoinAlgorithm.INLJ /* expectedJoinAlgorithm */, + {b: 1} /* indexKeyPattern */); + assert.commandWorked(foreignColl.dropIndexes()); + + // Verify that a leading $match won't filter out a legitimate wildcard index. + assert.commandWorked(foreignColl.createIndex({'$**': 1}, {wildcardProjection: {b: 1, c: 1}})); + runTest(coll, + [ + {$match: {'c.d': 1}}, + {$lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "out"}} + ], + JoinAlgorithm.NLJ /* expectedJoinAlgorithm */); + assert.commandWorked(foreignColl.deleteOne(mkDoc)); assert.commandWorked(foreignColl.dropIndexes()); })(); @@ -641,12 +698,136 @@ function setLookupPushdownDisabled(value) { assert(unionColl.drop()); }()); +// Test which verifies that the right side of a classic $lookup is never lowered into SBE, even if +// the queries for the right side are eligible on their own to run in SBE. +(function verifyThatClassicLookupRightSideIsNeverLoweredIntoSBE() { + // If running with SBE fully enabled, verify that our $match is SBE compatible. Otherwise, + // verify that the same $match, when used as a $lookup sub-pipeline, will not be lowered + // into SBE. + const subPipeline = [{$match: {b: 2}}]; + if (checkSBEEnabled(db, ["featureFlagSbeFull"])) { + const subPipelineExplain = foreignColl.explain().aggregate(subPipeline); + assert(subPipelineExplain.hasOwnProperty("explainVersion"), subPipelineExplain); + assert.eq(subPipelineExplain["explainVersion"], "2", subPipelineExplain); + } else { + const pipeline = [{$lookup: {from: foreignCollName, pipeline: subPipeline, as: "result"}}]; + runTest(coll, pipeline, JoinAlgorithm.Classic /* expectedJoinAlgorithm */); + + // Run the pipeline enough times to generate a cache entry for the right side in the foreign + // collection. + coll.aggregate(pipeline).itcount(); + coll.aggregate(pipeline).itcount(); + + const cacheEntries = foreignColl.getPlanCache().list(); + assert.eq(cacheEntries.length, 1); + const cacheEntry = cacheEntries[0]; + + // The cached plan should be a classic plan. + assert(cacheEntry.hasOwnProperty("version"), cacheEntry); + assert.eq(cacheEntry.version, "1", cacheEntry); + assert(cacheEntry.hasOwnProperty("cachedPlan"), cacheEntry); + const cachedPlan = cacheEntry.cachedPlan; + + // The cached plan should not have slot based plan. Instead, it should be a FETCH + IXSCAN + // executed in the classic engine. + assert(!cachedPlan.hasOwnProperty("slots"), cacheEntry); + assert(cachedPlan.hasOwnProperty("stage"), cacheEntry); + + assert(planHasStage(db, cachedPlan, "FETCH"), cacheEntry); + assert(planHasStage(db, cachedPlan, "IXSCAN"), cacheEntry); + } +}()); + MongoRunner.stopMongod(conn); +// Verify that pipeline stages get pushed down according to the subset of SBE that is enabled. +(function verifyPushdownLogicSbePartiallyEnabled() { + const conn = MongoRunner.runMongod( + {setParameter: {featureFlagSBELookupPushdown: true, allowDiskUseByDefault: false}}); + const db = conn.getDB(name); + if (checkSBEEnabled(db, ["featureFlagSbeFull"])) { + jsTestLog("Skipping test case because SBE is fully enabled, but this test case assumes" + + " that it is not fully enabled"); + MongoRunner.stopMongod(conn); + return; + } + const coll = db[name]; + const foreignColl = db[foreignCollName]; + + assert.commandWorked(coll.insert({a: 1})); + assert.commandWorked(foreignColl.insert([{b: 1}, {b: 1}])); + + const lookupStage = { + $lookup: {from: foreignCollName, localField: "a", foreignField: "b", as: "result"} + }; + const groupStage = {$group: {_id: "$a", avg: {$avg: "$b"}}}; + let pipeline = [lookupStage]; + let explain = coll.explain().aggregate(pipeline); + + // We should have exactly one $lookup stage and no EQ_LOOKUP nodes. + assert.eq(0, getAggPlanStages(explain, "EQ_LOOKUP").length, explain); + assert.eq(1, getAggPlanStages(explain, "$lookup").length, explain); + + // Run a pipeline where the $group is eligible for push down, but the $lookup is not. + pipeline = [groupStage, lookupStage]; + explain = coll.explain().aggregate(pipeline); + + // We should have exactly one $lookup stage and no EQ_LOOKUP nodes. + assert.eq(0, getAggPlanStages(explain, "EQ_LOOKUP").length, explain); + assert.eq(1, getAggPlanStages(explain, "$lookup").length, explain); + + // We should have exactly one GROUP node and no $group stages. + assert.eq(1, getAggPlanStages(explain, "GROUP").length, explain); + assert.eq(0, getAggPlanStages(explain, "$group").length, explain); + + // Run a pipeline where only the first $group is eligible for push down, but the rest of the + // stages are not. + pipeline = [groupStage, lookupStage, groupStage, lookupStage]; + explain = coll.explain().aggregate(pipeline); + + // We should have two $lookup stages and no EQ_LOOKUP nodes. + assert.eq(0, getAggPlanStages(explain, "EQ_LOOKUP").length, explain); + assert.eq(2, getAggPlanStages(explain, "$lookup").length, explain); + + // We should have one GROUP node and one $group stage. + assert.eq(1, getAggPlanStages(explain, "GROUP").length, explain); + assert.eq(1, getAggPlanStages(explain, "$group").length, explain); + + function assertEngine(pipeline, engine) { + const explain = coll.explain().aggregate(pipeline); + assert(explain.hasOwnProperty("explainVersion"), explain); + assert.eq(explain.explainVersion, engine === "sbe" ? "2" : "1"); + } + + const matchStage = {$match: {a: 1}}; + + // $group on its own is SBE compatible. + assertEngine([groupStage], "sbe" /* engine */); + + // $group with $match is also SBE compatible. + assertEngine([matchStage, groupStage], "sbe" /* engine */); + + // A $lookup rejected during engine selection should inhibit SBE. + assertEngine([lookupStage], "classic" /* engine */); + assertEngine([matchStage, lookupStage], "classic" /* engine */); + assertEngine([matchStage, lookupStage, groupStage], "classic" /* engine */); + + // Constructing an index over the foreignField of 'lookupStage' will cause the $lookup to be + // pushed down. + assert.commandWorked(foreignColl.createIndex({b: 1})); + assertEngine([matchStage, lookupStage, groupStage], "sbe" /* engine */); + assert.commandWorked(foreignColl.dropIndex({b: 1})); + + // Though the $lookup will not run in SBE, a preceding $group should still let SBE be used. + assertEngine([matchStage, groupStage, lookupStage], "sbe" /* engine */); + MongoRunner.stopMongod(conn); +}()); + (function testHashJoinQueryKnobs() { // Create a new scope and start a new mongod so that the mongod-wide global state changes do not // affect subsequent tests if any. - const conn = MongoRunner.runMongod({setParameter: "featureFlagSBELookupPushdown=true"}); + const conn = MongoRunner.runMongod( + {setParameter: {featureFlagSBELookupPushdown: true, featureFlagSbeFull: true}}); const db = conn.getDB(name); const lcoll = db.query_knobs_local; const fcoll = db.query_knobs_foreign; @@ -822,8 +1003,13 @@ const st = new ShardingTest({ shards: 2, mongos: 1, other: { - shardOptions: - {setParameter: {featureFlagSBELookupPushdown: true, allowDiskUseByDefault: false}} + shardOptions: { + setParameter: { + featureFlagSBELookupPushdown: true, + featureFlagSbeFull: true, + allowDiskUseByDefault: false + } + } } }); db = st.s.getDB(name); diff --git a/jstests/noPassthrough/plan_cache_replan_group_lookup.js b/jstests/noPassthrough/plan_cache_replan_group_lookup.js index 4eb954eb93c..2f749227316 100644 --- a/jstests/noPassthrough/plan_cache_replan_group_lookup.js +++ b/jstests/noPassthrough/plan_cache_replan_group_lookup.js @@ -164,17 +164,29 @@ function dropLookupForeignColl() { } const lookupPushdownEnabled = checkSBEEnabled(db, ["featureFlagSBELookupPushdown"]); +const lookupPushdownNLJEnabled = + checkSBEEnabled(db, ["featureFlagSBELookupPushdown", "featureFlagSbeFull"]); function verifyCorrectLookupAlgorithmUsed(targetJoinAlgorithm, pipeline, aggOptions = {}) { if (!lookupPushdownEnabled) { return; } + + if (!lookupPushdownNLJEnabled && targetJoinAlgorithm === "NestedLoopJoin") { + targetJoinAlgorithm = "Classic"; + } + const explain = coll.explain().aggregate(pipeline, aggOptions); const eqLookupNodes = getAggPlanStages(explain, "EQ_LOOKUP"); - // Verify via explain that $lookup was lowered and appropriate $lookup algorithm was chosen. - assert.eq( - eqLookupNodes.length, 1, "expected at least one EQ_LOOKUP node; got " + tojson(explain)); - assert.eq(eqLookupNodes[0].strategy, targetJoinAlgorithm); + if (targetJoinAlgorithm === "Classic") { + assert.eq(eqLookupNodes.length, 0, "expected no EQ_LOOKUP nodes; got " + tojson(explain)); + } else { + // Verify via explain that $lookup was lowered and appropriate $lookup algorithm was chosen. + assert.eq(eqLookupNodes.length, + 1, + "expected at least one EQ_LOOKUP node; got " + tojson(explain)); + assert.eq(eqLookupNodes[0].strategy, targetJoinAlgorithm); + } } // NLJ. diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp index 4fc6234b35c..3169859560f 100644 --- a/src/mongo/db/pipeline/pipeline_d.cpp +++ b/src/mongo/db/pipeline/pipeline_d.cpp @@ -80,6 +80,7 @@ #include "mongo/db/query/plan_executor_factory.h" #include "mongo/db/query/plan_executor_impl.h" #include "mongo/db/query/plan_summary_stats.h" +#include "mongo/db/query/planner_analysis.h" #include "mongo/db/query/query_feature_flags_gen.h" #include "mongo/db/query/query_knobs_gen.h" #include "mongo/db/query/query_planner.h" @@ -168,6 +169,7 @@ std::vector<std::unique_ptr<InnerPipelineStageInterface>> extractSbeCompatibleSt internalQuerySlotBasedExecutionDisableLookupPushdown.load() || isMainCollectionSharded || collections.isAnySecondaryNamespaceAViewOrSharded(); + std::map<NamespaceString, SecondaryCollectionInfo> secondaryCollInfo; for (auto itr = sources.begin(); itr != sources.end();) { const bool isLastSource = itr->get() == sources.back().get(); @@ -195,10 +197,31 @@ std::vector<std::unique_ptr<InnerPipelineStageInterface>> extractSbeCompatibleSt // Note that 'lookupStage->sbeCompatible()' encodes whether the foreign collection is a // view. if (lookupStage->sbeCompatible()) { - stagesForPushdown.push_back( - std::make_unique<InnerPipelineStageImpl>(lookupStage, isLastSource)); - sources.erase(itr++); - continue; + // Fill out secondary collection information to assist in deciding whether we should + // push down any $lookup stages into SBE. + // TODO SERVER-67024: This should be removed once NLJ is re-enabled by default. + if (secondaryCollInfo.empty()) { + secondaryCollInfo = + fillOutSecondaryCollectionsInformation(expCtx->opCtx, collections, cq); + } + + auto [strategy, _] = QueryPlannerAnalysis::determineLookupStrategy( + lookupStage->getFromNs().toString(), + lookupStage->getForeignField()->fullPath(), + secondaryCollInfo, + cq->getExpCtx()->allowDiskUse, + cq->getCollator()); + + // While we do support executing NLJ in SBE, this join algorithm does not currently + // perform as well as executing $lookup in the classic engine, so fall back to + // classic unless 'gFeatureFlagSbeFull' is enabled. + if (strategy != EqLookupNode::LookupStrategy::kNestedLoopJoin || + feature_flags::gFeatureFlagSbeFull.isEnabledAndIgnoreFCV()) { + stagesForPushdown.push_back( + std::make_unique<InnerPipelineStageImpl>(lookupStage, isLastSource)); + sources.erase(itr++); + continue; + } } break; } diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index 35c83e8c5a5..526244649a8 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -1389,9 +1389,6 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor( if (extractAndAttachPipelineStages) { extractAndAttachPipelineStages(canonicalQuery.get()); } - - // TODO SERVER-65960: Optionally refactor this logic once we have a mechanism to reattach - // pipeline stages. // Use SBE if we find any $group/$lookup stages eligible for execution in SBE or if SBE // is fully enabled. Otherwise, fallback to the classic engine. if (canonicalQuery->pipeline().empty() && diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp index 04f1e30104c..d410b8cec03 100644 --- a/src/mongo/db/query/planner_analysis.cpp +++ b/src/mongo/db/query/planner_analysis.cpp @@ -582,18 +582,6 @@ void removeProjectSimpleBelowGroupRecursive(QuerySolutionNode* solnRoot) { } } } -} // namespace - -// static -std::unique_ptr<QuerySolution> QueryPlannerAnalysis::removeProjectSimpleBelowGroup( - std::unique_ptr<QuerySolution> soln) { - auto root = soln->extractRoot(); - - removeProjectSimpleBelowGroupRecursive(root.get()); - - soln->setRoot(std::move(root)); - return soln; -} // Checks if the foreign collection is eligible for the hash join algorithm. We conservatively // choose the hash join algorithm for cases when the hash table is unlikely to spill data. @@ -607,13 +595,36 @@ bool isEligibleForHashJoin(const SecondaryCollectionInfo& foreignCollInfo) { internalQueryCollectionMaxStorageSizeBytesToChooseHashJoin.load(); } +// Determines whether 'index' is eligible for executing the right side of a pushed down $lookup over +// 'foreignField'. +bool isIndexEligibleForRightSideOfLookupPushdown(const IndexEntry& index, + const CollatorInterface* collator, + const std::string& foreignField) { + return (index.type == INDEX_BTREE || index.type == INDEX_HASHED) && + index.keyPattern.firstElement().fieldName() == foreignField && !index.filterExpr && + !index.sparse && CollatorInterface::collatorsMatch(collator, index.collator); +} +} // namespace + +// static +std::unique_ptr<QuerySolution> QueryPlannerAnalysis::removeProjectSimpleBelowGroup( + std::unique_ptr<QuerySolution> soln) { + auto root = soln->extractRoot(); + + removeProjectSimpleBelowGroupRecursive(root.get()); + + soln->setRoot(std::move(root)); + return soln; +} + // static -void QueryPlannerAnalysis::determineLookupStrategy( - EqLookupNode* eqLookupNode, +std::pair<EqLookupNode::LookupStrategy, boost::optional<IndexEntry>> +QueryPlannerAnalysis::determineLookupStrategy( + const std::string& foreignCollName, + const std::string& foreignField, const std::map<NamespaceString, SecondaryCollectionInfo>& collectionsInfo, bool allowDiskUse, const CollatorInterface* collator) { - const auto& foreignCollName = eqLookupNode->foreignCollection; auto foreignCollItr = collectionsInfo.find(NamespaceString(foreignCollName)); tassert(5842600, str::stream() << "Expected collection info, but found none; target collection: " @@ -641,11 +652,7 @@ void QueryPlannerAnalysis::determineLookupStrategy( }); for (const auto& index : indexes) { - if ((index.type == INDEX_BTREE || index.type == INDEX_HASHED) && - index.keyPattern.firstElement().fieldName() == - eqLookupNode->joinFieldForeign.fullPath() && - !index.filterExpr && !index.sparse && - CollatorInterface::collatorsMatch(collator, index.collator)) { + if (isIndexEligibleForRightSideOfLookupPushdown(index, collator, foreignField)) { return index; } } @@ -654,14 +661,13 @@ void QueryPlannerAnalysis::determineLookupStrategy( }(); if (!foreignCollItr->second.exists) { - eqLookupNode->lookupStrategy = EqLookupNode::LookupStrategy::kNonExistentForeignCollection; + return {EqLookupNode::LookupStrategy::kNonExistentForeignCollection, boost::none}; } else if (foreignIndex) { - eqLookupNode->lookupStrategy = EqLookupNode::LookupStrategy::kIndexedLoopJoin; - eqLookupNode->idxEntry = foreignIndex; + return {EqLookupNode::LookupStrategy::kIndexedLoopJoin, std::move(foreignIndex)}; } else if (allowDiskUse && isEligibleForHashJoin(foreignCollItr->second)) { - eqLookupNode->lookupStrategy = EqLookupNode::LookupStrategy::kHashJoin; + return {EqLookupNode::LookupStrategy::kHashJoin, boost::none}; } else { - eqLookupNode->lookupStrategy = EqLookupNode::LookupStrategy::kNestedLoopJoin; + return {EqLookupNode::LookupStrategy::kNestedLoopJoin, boost::none}; } } diff --git a/src/mongo/db/query/planner_analysis.h b/src/mongo/db/query/planner_analysis.h index 9f4a9a598a8..d7473336384 100644 --- a/src/mongo/db/query/planner_analysis.h +++ b/src/mongo/db/query/planner_analysis.h @@ -124,16 +124,20 @@ public: std::unique_ptr<QuerySolution> soln); /** - * For the provided 'eqLookupNode', determines what join algorithm should be used to execute it - * and marks the node accordingly. In particular: + * For the provided 'foreignCollName' and 'foreignFieldName' corresponding to an EqLookupNode, + * returns what join algorithm should be used to execute it. In particular: + * - An empty array is produced for each document if the foreign collection does not exist. * - An indexed nested loop join is chosen if an index on the foreign collection can be used to - * answer the join predicate. + * answer the join predicate. Also returns which index on the foreign collection should be + * used to answer the predicate. * - A hash join is chosen if disk use is allowed and if the foreign collection is sufficiently * small. * - A nested loop join is chosen in all other cases. */ - static void determineLookupStrategy( - EqLookupNode* eqLookupNode, + static std::pair<EqLookupNode::LookupStrategy, boost::optional<IndexEntry>> + determineLookupStrategy( + const std::string& foreignCollName, + const std::string& foreignField, const std::map<NamespaceString, SecondaryCollectionInfo>& collectionsInfo, bool allowDiskUse, const CollatorInterface* collator); diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp index f7a36dd682a..693ce386cdf 100644 --- a/src/mongo/db/query/query_planner.cpp +++ b/src/mongo/db/query/query_planner.cpp @@ -1362,8 +1362,8 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan( } /** - * The 'query' might contain parts of aggregation pipeline. For now, we plan those separately - * and later attach the agg portion of the plan to the solution(s) for the "find" part of the query. + * The 'query' might contain parts of aggregation pipeline. For now, we plan those separately and + * later attach the agg portion of the plan to the solution(s) for the "find" part of the query. */ std::unique_ptr<QuerySolution> QueryPlanner::extendWithAggPipeline( const CanonicalQuery& query, @@ -1391,17 +1391,21 @@ std::unique_ptr<QuerySolution> QueryPlanner::extendWithAggPipeline( tassert(6369000, "This $lookup stage should be compatible with SBE", lookupStage->sbeCompatible()); + auto [strategy, idxEntry] = QueryPlannerAnalysis::determineLookupStrategy( + lookupStage->getFromNs().toString(), + lookupStage->getForeignField()->fullPath(), + secondaryCollInfos, + query.getExpCtx()->allowDiskUse, + query.getCollator()); auto eqLookupNode = std::make_unique<EqLookupNode>(std::move(solnForAgg), lookupStage->getFromNs().toString(), lookupStage->getLocalField()->fullPath(), lookupStage->getForeignField()->fullPath(), lookupStage->getAsField().fullPath(), + strategy, + std::move(idxEntry), innerStage->isLastSource() /* shouldProduceBson */); - QueryPlannerAnalysis::determineLookupStrategy(eqLookupNode.get(), - secondaryCollInfos, - query.getExpCtx()->allowDiskUse, - query.getCollator()); solnForAgg = std::move(eqLookupNode); continue; } diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp index 6fb1e24e79a..893fef833e0 100644 --- a/src/mongo/db/query/query_solution.cpp +++ b/src/mongo/db/query/query_solution.cpp @@ -1609,6 +1609,8 @@ std::unique_ptr<QuerySolutionNode> EqLookupNode::clone() const { joinFieldLocal, joinFieldForeign, joinField, + lookupStrategy, + idxEntry, shouldProduceBson); return copy; } diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h index 34fe0566220..455c5aabcaa 100644 --- a/src/mongo/db/query/query_solution.h +++ b/src/mongo/db/query/query_solution.h @@ -1478,12 +1478,16 @@ struct EqLookupNode : public QuerySolutionNode { const FieldPath& joinFieldLocal, const FieldPath& joinFieldForeign, const FieldPath& joinField, + EqLookupNode::LookupStrategy lookupStrategy, + boost::optional<IndexEntry> idxEntry, bool shouldProduceBson) : QuerySolutionNode(std::move(child)), foreignCollection(foreignCollection), joinFieldLocal(joinFieldLocal), joinFieldForeign(joinFieldForeign), joinField(joinField), + lookupStrategy(lookupStrategy), + idxEntry(std::move(idxEntry)), shouldProduceBson(shouldProduceBson) {} StageType getType() const override { diff --git a/src/mongo/db/query/query_solution_test.cpp b/src/mongo/db/query/query_solution_test.cpp index bdba8b50f6b..55003d3f8c1 100644 --- a/src/mongo/db/query/query_solution_test.cpp +++ b/src/mongo/db/query/query_solution_test.cpp @@ -1133,7 +1133,14 @@ TEST(QuerySolutionTest, EqLookupNodeWithIndexScan) { scanNode->bounds.startKey = BSON("a" << 1 << "b" << 1); scanNode->bounds.endKey = BSON("a" << 1 << "b" << 1); - EqLookupNode node(std::move(scanNode), "col", "local", "foreign", "as", false); + EqLookupNode node(std::move(scanNode), + "col", + "local", + "foreign", + "as", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); node.computeProperties(); @@ -1161,7 +1168,14 @@ TEST(QuerySolutionTest, EqLookupNodeWithIndexScanFieldOverwrite) { scanNode->bounds.endKey = BSON("a" << 1 << "b" << 1 << "c" << "1"); - EqLookupNode node(std::move(scanNode), "col", "local", "foreign", "b", false); + EqLookupNode node(std::move(scanNode), + "col", + "local", + "foreign", + "b", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); node.computeProperties(); @@ -1229,8 +1243,15 @@ TEST(QuerySolutionTest, GetSecondaryNamespaceVectorOverSingleEqLookupNode) { auto scanNode = std::make_unique<IndexScanNode>(buildSimpleIndexEntry(BSON("a" << 1))); const NamespaceString mainNss("db.main"); const auto foreignColl = "db.col"; - auto root = std::make_unique<EqLookupNode>( - std::move(scanNode), foreignColl, "local", "remote", "b", false); + auto root = std::make_unique<EqLookupNode>(std::move(scanNode), + foreignColl, + "local", + "remote", + "b", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); + QuerySolution qs; qs.setRoot(std::move(root)); @@ -1243,8 +1264,15 @@ TEST(QuerySolutionTest, GetSecondaryNamespaceVectorOverSingleEqLookupNode) { TEST(QuerySolutionTest, GetSecondaryNamespaceVectorDeduplicatesMainNss) { auto scanNode = std::make_unique<IndexScanNode>(buildSimpleIndexEntry(BSON("a" << 1))); const NamespaceString mainNss("db.main"); - auto root = std::make_unique<EqLookupNode>( - std::move(scanNode), mainNss.toString(), "local", "remote", "b", false); + auto root = std::make_unique<EqLookupNode>(std::move(scanNode), + mainNss.toString(), + "local", + "remote", + "b", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); + QuerySolution qs; qs.setRoot(std::move(root)); @@ -1260,10 +1288,25 @@ TEST(QuerySolutionTest, GetSecondaryNamespaceVectorOverNestedEqLookupNodes) { const NamespaceString mainNss("db.main"); const auto foreignCollOne = "db.col"; const auto foreignCollTwo = "db.foo"; - auto childEqLookupNode = std::make_unique<EqLookupNode>( - std::move(scanNode), foreignCollOne, "local", "remote", "b", false); - auto parentEqLookupNode = std::make_unique<EqLookupNode>( - std::move(childEqLookupNode), foreignCollTwo, "local", "remote", "b", false); + auto childEqLookupNode = + std::make_unique<EqLookupNode>(std::move(scanNode), + foreignCollOne, + "local", + "remote", + "b", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); + + auto parentEqLookupNode = + std::make_unique<EqLookupNode>(std::move(childEqLookupNode), + foreignCollTwo, + "local", + "remote", + "b", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); QuerySolution qs; qs.setRoot(std::move(parentEqLookupNode)); @@ -1280,10 +1323,25 @@ TEST(QuerySolutionTest, GetSecondaryNamespaceVectorDeduplicatesNestedEqLookupNod auto scanNode = std::make_unique<IndexScanNode>(buildSimpleIndexEntry(BSON("a" << 1))); const NamespaceString mainNss("db.main"); const auto foreignColl = "db.col"; - auto childEqLookupNode = std::make_unique<EqLookupNode>( - std::move(scanNode), foreignColl, "local", "remote", "b", false); - auto parentEqLookupNode = std::make_unique<EqLookupNode>( - std::move(childEqLookupNode), foreignColl, "local", "remote", "b", false); + auto childEqLookupNode = + std::make_unique<EqLookupNode>(std::move(scanNode), + foreignColl, + "local", + "remote", + "b", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); + + auto parentEqLookupNode = + std::make_unique<EqLookupNode>(std::move(childEqLookupNode), + foreignColl, + "local", + "remote", + "b", + EqLookupNode::LookupStrategy::kNestedLoopJoin, + boost::none /* idxEntry */, + false /* shouldProduceBson */); QuerySolution qs; qs.setRoot(std::move(parentEqLookupNode)); diff --git a/src/mongo/db/query/sbe_runtime_planner.h b/src/mongo/db/query/sbe_runtime_planner.h index 7d28e268ae3..9078f79e70d 100644 --- a/src/mongo/db/query/sbe_runtime_planner.h +++ b/src/mongo/db/query/sbe_runtime_planner.h @@ -45,13 +45,13 @@ namespace mongo::sbe { * index pointing to the winning plan within this vector. */ struct CandidatePlans { - std::vector<plan_ranker::CandidatePlan> plans; - size_t winnerIdx; - auto& winner() { invariant(winnerIdx < plans.size()); return plans[winnerIdx]; } + + std::vector<plan_ranker::CandidatePlan> plans; + size_t winnerIdx; }; /** 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 281fa4c86d5..36e7fee3c09 100644 --- a/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp +++ b/src/mongo/db/query/sbe_stage_builder_lookup_test.cpp @@ -120,9 +120,14 @@ public: // Construct logical query solution. auto foreignCollName = _foreignNss.toString(); - auto lookupNode = std::make_unique<EqLookupNode>( - std::move(localScanNode), foreignCollName, localKey, foreignKey, asKey, true); - lookupNode->lookupStrategy = strategy; + auto lookupNode = std::make_unique<EqLookupNode>(std::move(localScanNode), + foreignCollName, + localKey, + foreignKey, + asKey, + strategy, + boost::none /* idxEntry */, + true /* shouldProduceBson */); auto solution = makeQuerySolution(std::move(lookupNode)); // Convert logical solution into the physical SBE plan. |