SERVER-115104 Pretty-print join plans as trees in e2e tests (#45048)

GitOrigin-RevId: f6c8ec2d9295e4f92a1e36785487ef9699fc03e5

jstests/libs/query/analyze_plan.js

@@ -171,7 +171,7 @@ function isPlainObject(value) {
     return value && typeof value == "object" && value.constructor === Object;
 }
 
-const kExplainChildFieldNames = [
+export const kExplainChildFieldNames = [
     "inputStage",
     "inputStages",
     "thenStage",

jstests/query_golden/join_opt/base_coll_reorder_md.js

@@ -5,10 +5,10 @@
  *   requires_fcv_83
  * ]
  */
-import {codeOneLine, linebreak, section, subSection} from "jstests/libs/pretty_md.js";
+import {linebreak, section, subSection} from "jstests/libs/pretty_md.js";
 import {outputAggregationPlanAndResults} from "jstests/libs/query/golden_test_utils.js";
-import {normalizePlan, getWinningPlanFromExplain} from "jstests/libs/query/analyze_plan.js";
 import {checkSbeFullFeatureFlagEnabled} from "jstests/libs/query/sbe_util.js";
+import {prettyPrintWinningPlan, getWinningJoinOrderOneLine} from "jstests/query_golden/libs/pretty_plan.js";
 
 const coll = db[jsTestName() + "_base"];
 coll.drop();
@@ -55,49 +55,11 @@ const origParams = assert.commandWorked(
 );
 delete origParams.ok;
 
-function getStageAbbreviation(stageName) {
-    switch (stageName) {
-        case "HASH_JOIN_EMBEDDING":
-            return "HJ";
-        case "NESTED_LOOP_JOIN":
-            return "NLJ";
-        case "MERGE_JOIN":
-            return "MJ";
-        default:
-            return stageName;
-    }
-}
-
-function formatEmbeddingField(field) {
-    if (field && field !== "none") {
-        return field;
-    }
-
-    return "_";
-}
-
-function abbreviate(node) {
-    const abbrev = getStageAbbreviation(node.stage);
-    if (abbrev == node.stage) {
-        return abbrev;
-    }
-    const l = formatEmbeddingField(node.leftEmbeddingField);
-    const r = formatEmbeddingField(node.rightEmbeddingField);
-    const children = node.inputStages.map(abbreviate);
-    assert.eq(children.length, 2);
-    return `(${abbrev} ${l} = ${children[0]}, ${r} = ${children[1]})`;
-}
-
-function getJoinOrder(explain) {
-    const winningPlan = normalizePlan(getWinningPlanFromExplain(explain), false /*shouldPlatten*/);
-    const x = abbreviate(winningPlan);
-    return x;
-}
-
 function runSingleTest(subtitle, pipeline, seen = undefined) {
     let joinOrder = undefined;
+    const explain = coll.explain().aggregate(pipeline);
     if (seen) {
-        joinOrder = getJoinOrder(coll.explain().aggregate(pipeline));
+        joinOrder = getWinningJoinOrderOneLine(explain);
         if (seen.has(joinOrder)) {
             return undefined;
         }
@@ -105,7 +67,7 @@ function runSingleTest(subtitle, pipeline, seen = undefined) {
     }
     subSection(subtitle);
     if (joinOrder) {
-        codeOneLine(joinOrder, true);
+        prettyPrintWinningPlan(explain);
     } else {
         outputAggregationPlanAndResults(coll, pipeline, {}, true, false);
     }

jstests/query_golden/join_opt/basic_joins_md.js

@@ -5,10 +5,12 @@
  *   requires_fcv_83
  * ]
  */
+import {assertArrayEq} from "jstests/aggregation/extras/utils.js";
 import {line, linebreak, section, subSection} from "jstests/libs/pretty_md.js";
 import {outputAggregationPlanAndResults} from "jstests/libs/query/golden_test_utils.js";
 import {getQueryPlanner} from "jstests/libs/query/analyze_plan.js";
 import {checkSbeFullFeatureFlagEnabled} from "jstests/libs/query/sbe_util.js";
+import {prettyPrintWinningPlan} from "jstests/query_golden/libs/pretty_plan.js";
 
 const coll = db[jsTestName()];
 coll.drop();
@@ -39,8 +41,7 @@ assert.commandWorked(
     ]),
 );
 
-function verifyExplainOutput(pipeline, joinOptExpectedInExplainOutput) {
-    const explain = coll.explain().aggregate(pipeline);
+function verifyExplainOutput(explain, joinOptExpectedInExplainOutput) {
     const winningPlan = getQueryPlanner(explain).winningPlan;
 
     if (joinOptExpectedInExplainOutput) {
@@ -55,131 +56,100 @@ function verifyExplainOutput(pipeline, joinOptExpectedInExplainOutput) {
     assert(!("usedJoinOptimization" in winningPlan), winningPlan);
 }
 
+function getJoinTestResultsAndExplain(desc, pipeline, params) {
+    subSection(desc);
+    assert.commandWorked(db.adminCommand({setParameter: 1, ...params}));
+    return [coll.aggregate(pipeline).toArray(), coll.explain().aggregate(pipeline)];
+}
+
+function runJoinTestAndCompare(desc, pipeline, params, expected) {
+    const [actual, explain] = getJoinTestResultsAndExplain(desc, pipeline, params);
+    assertArrayEq({expected, actual});
+    verifyExplainOutput(explain, true /* joinOptExpectedInExplainOutput */);
+    prettyPrintWinningPlan(explain);
+}
+
 function runBasicJoinTest(pipeline) {
     try {
         subSection("No join opt");
         assert.commandWorked(db.adminCommand({setParameter: 1, internalEnableJoinOptimization: false}));
         outputAggregationPlanAndResults(coll, pipeline, {}, true, false, false /* noLineBreak*/);
+        const noJoinExplain = coll.explain().aggregate(pipeline);
         const noJoinOptResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, false /* joinOptExpectedInExplainOutput */);
+        verifyExplainOutput(noJoinExplain, false /* joinOptExpectedInExplainOutput */);
 
-        subSection("With bottom-up plan enumeration (left-deep)");
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalEnableJoinOptimization: true}));
-        assert.commandWorked(
-            db.adminCommand({
-                setParameter: 1,
+        runJoinTestAndCompare(
+            "With bottom-up plan enumeration (left-deep)",
+            pipeline,
+            {
+                internalEnableJoinOptimization: true,
                 internalJoinReorderMode: "bottomUp",
                 internalJoinPlanTreeShape: "leftDeep",
-            }),
-        );
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        const bottomUpLeftDeepResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, bottomUpLeftDeepResults),
-            "Results differ between no join opt and bottom-up left-deep join enumeration",
+            },
+            noJoinOptResults,
         );
 
-        subSection("With bottom-up plan enumeration (right-deep)");
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalJoinPlanTreeShape: "rightDeep"}));
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        const bottomUpRightDeepResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, bottomUpRightDeepResults),
-            "Results differ between no join opt and bottom-up right-deep join enumeration",
+        runJoinTestAndCompare(
+            "With bottom-up plan enumeration (right-deep)",
+            pipeline,
+            {internalJoinPlanTreeShape: "rightDeep"},
+            noJoinOptResults,
         );
 
-        subSection("With bottom-up plan enumeration (zig-zag)");
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalJoinPlanTreeShape: "zigZag"}));
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        const bottomUpZigZagResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, bottomUpZigZagResults),
-            "Results differ between no join opt and bottom-up zig-zag join enumeration",
+        runJoinTestAndCompare(
+            "With bottom-up plan enumeration (zig-zag)",
+            pipeline,
+            {internalJoinPlanTreeShape: "zigZag"},
+            noJoinOptResults,
         );
 
-        subSection("With random order, seed 44, nested loop joins");
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalJoinReorderMode: "random"}));
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalRandomJoinOrderSeed: 44}));
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        const seed44NLJResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, seed44NLJResults),
-            "Results differ between no join opt and seed 44 NLJ",
+        for (const internalRandomJoinOrderSeed of [44, 45]) {
+            runJoinTestAndCompare(
+                `With random order, seed ${internalRandomJoinOrderSeed}, nested loop joins`,
+                pipeline,
+                {internalJoinReorderMode: "random", internalRandomJoinOrderSeed},
+                noJoinOptResults,
+            );
+
+            runJoinTestAndCompare(
+                `With random order, seed ${internalRandomJoinOrderSeed}, hash join enabled`,
+                pipeline,
+                {internalRandomJoinReorderDefaultToHashJoin: true},
+                noJoinOptResults,
+            );
+        }
+
+        // Run tests with indexes.
+        assert.commandWorked(foreignColl1.createIndex({a: 1}));
+        assert.commandWorked(foreignColl2.createIndex({b: 1}));
+
+        runJoinTestAndCompare(
+            "With fixed order, index join",
+            pipeline,
+            {internalRandomJoinReorderDefaultToHashJoin: false},
+            noJoinOptResults,
         );
 
-        subSection("With random order, seed 44, hash join enabled");
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalRandomJoinReorderDefaultToHashJoin: true}));
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        const seed44HJResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, seed44HJResults),
-            "Results differ between no join opt and seed 44 HJ",
+        runJoinTestAndCompare(
+            "With bottom-up plan enumeration and indexes",
+            pipeline,
+            {internalJoinReorderMode: "bottomUp", internalJoinPlanTreeShape: "leftDeep"},
+            noJoinOptResults,
         );
 
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalRandomJoinReorderDefaultToHashJoin: false}));
-
-        subSection("With random order, seed 420, nested loop joins");
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalEnableJoinOptimization: true}));
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalRandomJoinOrderSeed: 420}));
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        const seed420NLJResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, seed420NLJResults),
-            "Results differ between no join opt and seed 420 NLJ",
-        );
-
-        subSection("With random order, seed 420, hash join enabled");
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalRandomJoinReorderDefaultToHashJoin: true}));
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        const seed420HJResults = coll.aggregate(pipeline).toArray();
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, seed420HJResults),
-            "Results differ between no join opt and seed 420 HJ",
-        );
-
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalRandomJoinReorderDefaultToHashJoin: false}));
-        foreignColl1.createIndex({a: 1});
-        foreignColl2.createIndex({b: 1});
-        subSection("With fixed order, index join");
-
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false, true /* noLineBreak*/);
-        verifyExplainOutput(pipeline, true /* joinOptExpectedInExplainOutput */);
-        const seedINLJResults = coll.aggregate(pipeline).toArray();
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, seedINLJResults),
-            "Results differ between no join opt and INLJ",
-        );
-
-        subSection("With bottom-up plan enumeration and indexes");
-        assert.commandWorked(
-            db.adminCommand({
-                setParameter: 1,
-                internalJoinReorderMode: "bottomUp",
-                internalJoinPlanTreeShape: "leftDeep",
-            }),
-        );
-        outputAggregationPlanAndResults(coll, pipeline, {}, true, false);
-        const bottomUpINLJResults = coll.aggregate(pipeline).toArray();
-        assert(
-            _resultSetsEqualUnordered(noJoinOptResults, bottomUpINLJResults),
-            "Results differ between no join opt and INLJ",
-        );
-
-        foreignColl1.dropIndex({a: 1});
-        foreignColl2.dropIndex({b: 1});
+        assert.commandWorked(foreignColl1.dropIndex({a: 1}));
+        assert.commandWorked(foreignColl2.dropIndex({b: 1}));
     } finally {
         // Reset flags.
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalEnableJoinOptimization: false}));
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalRandomJoinReorderDefaultToHashJoin: false}));
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalJoinReorderMode: "bottomUp"}));
-        assert.commandWorked(db.adminCommand({setParameter: 1, internalJoinPlanTreeShape: "zigZag"}));
+        assert.commandWorked(
+            db.adminCommand({
+                setParameter: 1,
+                internalEnableJoinOptimization: false,
+                internalRandomJoinReorderDefaultToHashJoin: false,
+                internalJoinReorderMode: "bottomUp",
+                internalJoinPlanTreeShape: "zigZag",
+            }),
+        );
     }
 }
 

jstests/query_golden/libs/pretty_plan.js

@@ -0,0 +1,123 @@
+/**
+ * Pretty-printing helpers for query plans in explain.
+ */
+import {prettyPrintTree} from "jstests/query_golden/libs/pretty_tree.js";
+import {normalizePlan, getWinningPlanFromExplain, kExplainChildFieldNames} from "jstests/libs/query/analyze_plan.js";
+
+//
+// Helpers used below to get an abbreviated join order.
+//
+
+function getStageAbbreviation(stageName) {
+    switch (stageName) {
+        case "HASH_JOIN_EMBEDDING":
+            return "HJ";
+        case "NESTED_LOOP_JOIN":
+            return "NLJ";
+        case "INDEX_NESTED_LOOP_JOIN":
+            return "INLJ";
+        default:
+            return stageName;
+    }
+}
+
+function formatEmbeddingField(field) {
+    if (field && field !== "none") {
+        return field;
+    }
+    return "_";
+}
+
+function abbreviate(node) {
+    const abbrev = getStageAbbreviation(node.stage);
+    if (abbrev == node.stage) {
+        if (node.nss) {
+            return `${node.stage} [${node.nss}]`;
+        }
+        return abbrev;
+    }
+    const l = formatEmbeddingField(node.leftEmbeddingField);
+    const r = formatEmbeddingField(node.rightEmbeddingField);
+    const children = node.inputStages.map(abbreviate);
+    assert.eq(children.length, 2);
+    return `(${abbrev} ${l} = ${children[0]}, ${r} = ${children[1]})`;
+}
+
+//
+// End of helpers.
+//
+
+/**
+ * Helper function to extract a one line join order from the input plan.
+ * Useful for checking if this is a duplicate join order.
+ */
+export function getJoinOrderOneLine(plan) {
+    const winningPlan = normalizePlan(plan, false /*shouldFlatten*/);
+    const x = abbreviate(winningPlan);
+    return x;
+}
+
+/**
+ * Same as above, but for the winning plan in an 'explain'.
+ */
+export function getWinningJoinOrderOneLine(explain) {
+    return getJoinOrderOneLine(getWinningPlanFromExplain(explain));
+}
+
+//
+// Helpers used below to get a query plan pretty-printed as a tree.
+//
+
+function shortenField(node, fieldName) {
+    return node[fieldName] ? " [" + node[fieldName] + "]" : "";
+}
+
+function printPlanNode(node) {
+    // Place any fields with special formatting/ that we don't want in the default output here.
+    const bannedFieldNames = ["stage", "planNodeId", "nss", "joinPredicates"].concat(kExplainChildFieldNames);
+
+    const entries = Object.entries(node).filter(([f, _]) => !bannedFieldNames.includes(f));
+    let str = `${node.stage}${shortenField(node, "nss")}${shortenField(node, "joinPredicates")}\n`;
+    for (let i = 0; i < entries.length; i++) {
+        const [f, v] = entries[i];
+        if (f == "filter" && Object.entries(v).length == 0) {
+            // Omit empty filters.
+            continue;
+        }
+        str += `${f}: ${tojsononeline(v)}${i < entries.length - 1 ? "\n" : ""}`;
+    }
+    return str;
+}
+
+function getNodeChildren(node) {
+    let children = [];
+    for (const [field, entry] of Object.entries(node)) {
+        if (kExplainChildFieldNames.includes(field)) {
+            if (Array.isArray(entry)) {
+                children = children.concat(entry);
+            } else {
+                children.push(entry);
+            }
+        }
+    }
+    return children;
+}
+
+//
+// End of helpers.
+//
+
+/**
+ * Pretty-prints the given plan as a tree.
+ */
+export function prettyPrintPlan(plan) {
+    const winningPlan = normalizePlan(plan, false /*shouldFlatten*/);
+    prettyPrintTree(winningPlan, printPlanNode, getNodeChildren);
+}
+
+/**
+ * Same as above, but extracts the winning plan from 'explain' to print.
+ */
+export function prettyPrintWinningPlan(explain) {
+    return prettyPrintPlan(getWinningPlanFromExplain(explain));
+}

jstests/query_golden/libs/pretty_tree.js

@@ -0,0 +1,68 @@
+/**
+ * Generic tree-printer for golden tests.
+ */
+import {code} from "jstests/libs/pretty_md.js";
+
+function defaultNodeToString(node) {
+    return tojsononeline(node);
+}
+
+function defaultGetChildrenForNode(node) {
+    return node.children || [];
+}
+
+const kWhitespaceSeparator = "  ";
+const kEdge = "|";
+function prefix(numEdges) {
+    return (kWhitespaceSeparator + kEdge).repeat(numEdges + 1);
+}
+
+function formatMultilineNodeStr(nodeStr, numEdges) {
+    const lines = nodeStr.split("\n");
+    const edges = "\n" + prefix(numEdges - 1) + kWhitespaceSeparator;
+    const prevEdges = prefix(numEdges) + "\n";
+    return prevEdges + prefix(numEdges - 1) + kWhitespaceSeparator + lines.join(edges);
+}
+
+function formatNodeForTree(nodeStr, depth, numEdges) {
+    if (depth == 0) {
+        // Root node.
+        return nodeStr + "\n";
+    }
+
+    return formatMultilineNodeStr(nodeStr, numEdges) + "\n";
+}
+
+/**
+ * Main entry point for generating a pretty-printed tree string.
+ */
+export function prettyTreeString(
+    node,
+    nodeLambda = defaultNodeToString,
+    getChildrenForNode = defaultGetChildrenForNode,
+    depth = 0,
+    numEdges = 0,
+) {
+    // Print root.
+    const printedNode = nodeLambda(node);
+    let out = formatNodeForTree(printedNode, depth, numEdges);
+
+    // Traverse tree- we print the last child first (highest in the tree).
+    const children = getChildrenForNode(node);
+    for (let i = children.length - 1; i >= 0; i--) {
+        // Increase distance from left & number of edges to print by index of child.
+        out += prettyTreeString(children[i], nodeLambda, getChildrenForNode, depth + 1, numEdges + i);
+    }
+    return out;
+}
+
+/**
+ * Main entry point for markdown-formatted tree string.
+ */
+export function prettyPrintTree(
+    node,
+    nodeLambda = defaultNodeToString,
+    getChildrenForNode = defaultGetChildrenForNode,
+) {
+    code(prettyTreeString(node, nodeLambda, getChildrenForNode), "");
+}