SERVER-64614 Update lookup_with_limit.js and look_with_limit_sharded.js for lowered $lookup plan

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"]);
 }());