From 22f72dd4701f6a6a0d6cda810a4fde4cb36bf895 Mon Sep 17 00:00:00 2001 From: Rui Liu Date: Sat, 2 Apr 2022 11:28:56 +0000 Subject: SERVER-64614 Update lookup_with_limit.js and look_with_limit_sharded.js for lowered $lookup plan --- jstests/libs/analyze_plan.js | 13 ++++++++ jstests/noPassthrough/lookup_with_limit_sharded.js | 36 +++++++++++++--------- .../noPassthroughWithMongod/lookup_with_limit.js | 34 ++++++++++---------- 3 files changed, 52 insertions(+), 31 deletions(-) diff --git a/jstests/libs/analyze_plan.js b/jstests/libs/analyze_plan.js index 24624c5ad59..dcfb3f11221 100644 --- a/jstests/libs/analyze_plan.js +++ b/jstests/libs/analyze_plan.js @@ -532,3 +532,16 @@ function getPlanCacheKeyFromShape({query = {}, projection = {}, sort = {}, colle return getPlanCacheKeyFromExplain(explainRes, db); } + +/** + * Given the winning query plan, flatten query plan tree into a list of plan stage names. + */ +function flattenQueryPlanTree(winningPlan) { + let stages = []; + while (winningPlan) { + stages.push(winningPlan.stage); + winningPlan = winningPlan.inputStage; + } + stages.reverse(); + return stages; +} diff --git a/jstests/noPassthrough/lookup_with_limit_sharded.js b/jstests/noPassthrough/lookup_with_limit_sharded.js index 19011f6d807..cb4d6953d1e 100644 --- a/jstests/noPassthrough/lookup_with_limit_sharded.js +++ b/jstests/noPassthrough/lookup_with_limit_sharded.js @@ -19,9 +19,8 @@ load("jstests/libs/sbe_util.js"); // For checkSBEEnabled. const st = new ShardingTest({shards: 2, config: 1}); const db = st.s.getDB("test"); -// TODO SERVER-64614 Update test cases for the SBE $lookup and remove 'if' block. -if (checkSBEEnabled(db, ["featureFlagSBELookupPushdown"])) { - jsTestLog("Skipping test because SBE and SBE $lookup features are both enabled."); +if (!checkSBEEnabled(db, ["featureFlagSBELookupPushdown"])) { + jsTestLog("Skipping test because SBE $lookup is not enabled."); st.stop(); return; } @@ -31,14 +30,20 @@ const other = db.lookup_with_limit_other; coll.drop(); other.drop(); -// Checks that the order of the pipeline stages matches the expected optimized ordering for an -// unsharded collection. -function checkUnshardedResults(pipeline, expectedPlanStage, expectedPipeline) { +// Checks that the order of the query stages and pipeline stages matches the expected optimized +// ordering for an unsharded collection. +function checkUnshardedResults(pipeline, expectedPlanStages, expectedPipeline) { const explain = coll.explain().aggregate(pipeline); - assert.eq( - getWinningPlan(explain.stages[0].$cursor.queryPlanner).stage, expectedPlanStage, explain); - for (let i = 0; i < expectedPipeline.length; i++) { - assert.eq(Object.keys(explain.stages[i + 1]), expectedPipeline[i], explain); + if (explain.stages) { + const queryStages = + flattenQueryPlanTree(getWinningPlan(explain.stages[0].$cursor.queryPlanner)); + const pipelineStages = explain.stages.slice(1).map(s => Object.keys(s)[0]); + assert.eq(queryStages, expectedPlanStages, explain); + assert.eq(pipelineStages, expectedPipeline, explain); + } else { + const queryStages = flattenQueryPlanTree(getWinningPlan(explain.queryPlanner)); + assert.eq(queryStages, expectedPlanStages, explain); + assert.eq([], expectedPipeline, explain); } } @@ -67,7 +72,7 @@ const lookupPipeline = [ {$lookup: {from: other.getName(), localField: "x", foreignField: "x", as: "from_other"}}, {$limit: 5} ]; -checkUnshardedResults(lookupPipeline, "LIMIT", ["$lookup"]); +checkUnshardedResults(lookupPipeline, ["COLLSCAN", "LIMIT", "EQ_LOOKUP"], []); // Check that lookup->addFields->lookup->limit is reordered to limit->lookup->addFields->lookup, // with the limit stage pushed down to query system. @@ -77,7 +82,8 @@ const multiLookupPipeline = [ {$lookup: {from: other.getName(), localField: "x", foreignField: "x", as: "additional"}}, {$limit: 5} ]; -checkUnshardedResults(multiLookupPipeline, "LIMIT", ["$lookup", "$addFields", "$lookup"]); +checkUnshardedResults( + multiLookupPipeline, ["COLLSCAN", "LIMIT", "EQ_LOOKUP"], ["$addFields", "$lookup"]); // Check that lookup->unwind->limit is reordered to lookup->limit, with the unwind stage being // absorbed into the lookup stage and preventing the limit from swapping before it. @@ -86,7 +92,7 @@ const unwindPipeline = [ {$unwind: "$from_other"}, {$limit: 5} ]; -checkUnshardedResults(unwindPipeline, "COLLSCAN", ["$lookup", "$limit"]); +checkUnshardedResults(unwindPipeline, ["COLLSCAN"], ["$lookup", "$limit"]); // Check that lookup->unwind->sort->limit is reordered to lookup->sort, with the unwind stage being // absorbed into the lookup stage and preventing the limit from swapping before it, and the limit @@ -106,9 +112,9 @@ const topKSortPipeline = [ {$lookup: {from: other.getName(), localField: "x", foreignField: "x", as: "from_other"}}, {$limit: 5} ]; -checkUnshardedResults(topKSortPipeline, "SORT", ["$lookup"]); +checkUnshardedResults(topKSortPipeline, ["COLLSCAN", "SORT", "EQ_LOOKUP"], []); const explain = coll.explain().aggregate(topKSortPipeline); -assert.eq(getWinningPlan(explain.stages[0].$cursor.queryPlanner).limitAmount, 5, explain); +assert.eq(getPlanStage(getWinningPlan(explain.queryPlanner), "SORT").limitAmount, 5, explain); // Tests on a sharded collection. coll.createIndex({x: 1}); diff --git a/jstests/noPassthroughWithMongod/lookup_with_limit.js b/jstests/noPassthroughWithMongod/lookup_with_limit.js index f8c22473b59..28dbb4f4702 100644 --- a/jstests/noPassthroughWithMongod/lookup_with_limit.js +++ b/jstests/noPassthroughWithMongod/lookup_with_limit.js @@ -7,8 +7,8 @@ load('jstests/libs/analyze_plan.js'); // For getWinningPlan(). load("jstests/libs/sbe_util.js"); // For checkSBEEnabled. -if (checkSBEEnabled(db, ["featureFlagSBELookupPushdown"])) { - jsTestLog("Skipping test because SBE and SBE $lookup features are both enabled."); +if (!checkSBEEnabled(db, ["featureFlagSBELookupPushdown"])) { + jsTestLog("Skipping test because SBE $lookup is not enabled."); return; } @@ -17,20 +17,22 @@ const other = db.lookup_with_limit_other; coll.drop(); other.drop(); -// Checks that the order of the pipeline stages matches the expected ordering depending on whether -// the pipeline is optimized or not. +// Checks that the order of the query stages and pipeline stages matches the expected ordering +// depending on whether the pipeline is optimized or not. function checkResults(pipeline, isOptimized, expected) { assert.commandWorked(db.adminCommand({ "configureFailPoint": 'disablePipelineOptimization', "mode": isOptimized ? 'off' : 'alwaysOn' })); const explain = coll.explain().aggregate(pipeline); - if (expected.length > 0) { - assert.eq( - getWinningPlan(explain.stages[0].$cursor.queryPlanner).stage, expected[0], explain); - } - for (let i = 1; i < expected.length; i++) { - assert.eq(Object.keys(explain.stages[i]), expected[i], explain); + if (explain.stages) { + const queryStages = + flattenQueryPlanTree(getWinningPlan(explain.stages[0].$cursor.queryPlanner)); + const pipelineStages = explain.stages.slice(1).map(s => Object.keys(s)[0]); + assert.eq(queryStages.concat(pipelineStages), expected, explain); + } else { + const queryStages = flattenQueryPlanTree(getWinningPlan(explain.queryPlanner)); + assert.eq(queryStages, expected, explain); } } @@ -51,8 +53,8 @@ var pipeline = [ {$lookup: {from: other.getName(), localField: "x", foreignField: "x", as: "from_other"}}, {$limit: 5} ]; -checkResults(pipeline, false, ["COLLSCAN", "$lookup", "$limit"]); -checkResults(pipeline, true, ["LIMIT", "$lookup"]); +checkResults(pipeline, false, ["COLLSCAN", "EQ_LOOKUP", "$limit"]); +checkResults(pipeline, true, ["COLLSCAN", "LIMIT", "EQ_LOOKUP"]); // Check that lookup->addFields->lookup->limit is reordered to limit->lookup->addFields->lookup, // with the limit stage pushed down to query system. @@ -62,8 +64,8 @@ pipeline = [ {$lookup: {from: other.getName(), localField: "x", foreignField: "x", as: "additional"}}, {$limit: 5} ]; -checkResults(pipeline, false, ["COLLSCAN", "$lookup", "$addFields", "$lookup", "$limit"]); -checkResults(pipeline, true, ["LIMIT", "$lookup", "$addFields", "$lookup"]); +checkResults(pipeline, false, ["COLLSCAN", "EQ_LOOKUP", "$addFields", "$lookup", "$limit"]); +checkResults(pipeline, true, ["COLLSCAN", "LIMIT", "EQ_LOOKUP", "$addFields", "$lookup"]); // Check that lookup->unwind->limit is reordered to lookup->limit, with the unwind stage being // absorbed into the lookup stage and preventing the limit from swapping before it. @@ -72,7 +74,7 @@ pipeline = [ {$unwind: "$from_other"}, {$limit: 5} ]; -checkResults(pipeline, false, ["COLLSCAN", "$lookup", "$unwind", "$limit"]); +checkResults(pipeline, false, ["COLLSCAN", "EQ_LOOKUP", "$unwind", "$limit"]); checkResults(pipeline, true, ["COLLSCAN", "$lookup", "$limit"]); // Check that lookup->unwind->sort->limit is reordered to lookup->sort, with the unwind stage being @@ -84,6 +86,6 @@ pipeline = [ {$sort: {x: 1}}, {$limit: 5} ]; -checkResults(pipeline, false, ["COLLSCAN", "$lookup", "$unwind", "$sort", "$limit"]); +checkResults(pipeline, false, ["COLLSCAN", "EQ_LOOKUP", "$unwind", "$sort", "$limit"]); checkResults(pipeline, true, ["COLLSCAN", "$lookup", "$sort"]); }()); -- cgit v1.2.1