SERVER-65388 Pass PrefixId by reference in optimizeGroup()

6 files changed, 83 insertions, 96 deletions

diff --git a/src/mongo/db/query/optimizer/cascades/physical_rewriter.cpp b/src/mongo/db/query/optimizer/cascades/physical_rewriter.cpp
index 0f979dc8e1a..9be06e8b76b 100644
--- a/src/mongo/db/query/optimizer/cascades/physical_rewriter.cpp
+++ b/src/mongo/db/query/optimizer/cascades/physical_rewriter.cpp
@@ -138,7 +138,7 @@ void PhysicalRewriter::costAndRetainBestNode(ABT node,
                                              ChildPropsType childProps,
                                              NodeCEMap nodeCEMap,
                                              const GroupIdType groupId,
-                                             const PrefixId& prefixId,
+                                             PrefixId& prefixId,
                                              PhysOptimizationResult& bestResult) {
     const CostAndCE nodeCostAndCE =
         _costDerivation.deriveCost(_memo, bestResult._physProps, node.ref(), childProps, nodeCEMap);
@@ -182,7 +182,7 @@ void PhysicalRewriter::costAndRetainBestNode(ABT node,
  */
 std::pair<bool, CostType> PhysicalRewriter::optimizeChildren(const CostType nodeCost,
                                                              ChildPropsType childProps,
-                                                             const PrefixId& prefixId,
+                                                             PrefixId& prefixId,
                                                              const CostType costLimit) {
     const bool disableBranchAndBound = _hints._disableBranchAndBound;
 
@@ -226,7 +226,7 @@ PhysicalRewriter::OptimizeGroupResult::OptimizeGroupResult(const size_t index, c
 
 PhysicalRewriter::OptimizeGroupResult PhysicalRewriter::optimizeGroup(const GroupIdType groupId,
                                                                       PhysProps physProps,
-                                                                      PrefixId prefixId,
+                                                                      PrefixId& prefixId,
                                                                       CostType costLimit) {
     const size_t localPlanExplorationCount = ++_memo._stats._physPlanExplorationCount;
     if (_memo.getDebugInfo().hasDebugLevel(2)) {
diff --git a/src/mongo/db/query/optimizer/cascades/physical_rewriter.h b/src/mongo/db/query/optimizer/cascades/physical_rewriter.h
index b162bdd9c05..0e59e732845 100644
--- a/src/mongo/db/query/optimizer/cascades/physical_rewriter.h
+++ b/src/mongo/db/query/optimizer/cascades/physical_rewriter.h
@@ -65,7 +65,7 @@ public:
      */
     OptimizeGroupResult optimizeGroup(GroupIdType groupId,
                                       properties::PhysProps physProps,
-                                      PrefixId prefixId,
+                                      PrefixId& prefixId,
                                       CostType costLimit);
 
 private:
@@ -73,12 +73,12 @@ private:
                                ChildPropsType childProps,
                                NodeCEMap nodeCEMap,
                                GroupIdType groupId,
-                               const PrefixId& prefixId,
+                               PrefixId& prefixId,
                                PhysOptimizationResult& bestResult);
 
     std::pair<bool, CostType> optimizeChildren(CostType nodeCost,
                                                ChildPropsType childProps,
-                                               const PrefixId& prefixId,
+                                               PrefixId& prefixId,
                                                CostType costLimit);
 
     // We don't own any of this.
diff --git a/src/mongo/db/query/optimizer/explain.cpp b/src/mongo/db/query/optimizer/explain.cpp
index f98c016cf94..ca98ec6a4d0 100644
--- a/src/mongo/db/query/optimizer/explain.cpp
+++ b/src/mongo/db/query/optimizer/explain.cpp
@@ -37,14 +37,8 @@
 namespace mongo::optimizer {
 
 BSONObj ABTPrinter::explainBSON() const {
-    auto [tag, val] = optimizer::ExplainGenerator::explainBSON(
+    return ExplainGenerator::explainBSONObj(
         _abtTree, true /*displayProperties*/, nullptr /*Memo*/, _nodeToPropsMap);
-    uassert(6624070, "Expected an object", tag == sbe::value::TypeTags::Object);
-    sbe::value::ValueGuard vg(tag, val);
-
-    BSONObjBuilder builder;
-    sbe::bson::convertToBsonObj(builder, sbe::value::getObjectView(val));
-    return builder.done().getOwned();
 }
 
 enum class ExplainVersion { V1, V2, V3, Vmax };
@@ -2282,6 +2276,20 @@ std::pair<sbe::value::TypeTags, sbe::value::Value> ExplainGenerator::explainBSON
     return gen.generate(node).moveValue();
 }
 
+BSONObj ExplainGenerator::explainBSONObj(const ABT& node,
+                                         const bool displayProperties,
+                                         const cascades::Memo* memo,
+                                         const NodeToGroupPropsMap& nodeMap) {
+    auto [tag, val] =
+        optimizer::ExplainGenerator::explainBSON(node, displayProperties, memo, nodeMap);
+    uassert(6624070, "Expected an object", tag == sbe::value::TypeTags::Object);
+    sbe::value::ValueGuard vg(tag, val);
+
+    BSONObjBuilder builder;
+    sbe::bson::convertToBsonObj(builder, sbe::value::getObjectView(val));
+    return builder.done().getOwned();
+}
+
 template <class PrinterType>
 static void printBSONstr(PrinterType& printer,
                          const sbe::value::TypeTags tag,
diff --git a/src/mongo/db/query/optimizer/explain.h b/src/mongo/db/query/optimizer/explain.h
index 970d912cdbc..5c79006143b 100644
--- a/src/mongo/db/query/optimizer/explain.h
+++ b/src/mongo/db/query/optimizer/explain.h
@@ -81,6 +81,11 @@ public:
         const cascades::Memo* memo = nullptr,
         const NodeToGroupPropsMap& nodeMap = {});
 
+    static BSONObj explainBSONObj(const ABT& node,
+                                  bool displayProperties = false,
+                                  const cascades::Memo* memo = nullptr,
+                                  const NodeToGroupPropsMap& nodeMap = {});
+
     static std::string printBSON(sbe::value::TypeTags tag, sbe::value::Value val);
 
     static std::string explainLogicalProps(const std::string& description,
diff --git a/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp b/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp
index b8ff9d71a13..32245ee586b 100644
--- a/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp
+++ b/src/mongo/db/query/optimizer/physical_rewriter_optimizer_test.cpp
@@ -1525,37 +1525,37 @@ TEST(PhysRewriter, FilterIndexingStress) {
         "|   |   Const [true]\n"
         "|   Filter []\n"
         "|   |   EvalFilter []\n"
-        "|   |   |   Variable [evalTemp_17]\n"
+        "|   |   |   Variable [evalTemp_81]\n"
         "|   |   PathTraverse []\n"
         "|   |   PathCompare [Eq]\n"
         "|   |   Const [0]\n"
         "|   Filter []\n"
         "|   |   EvalFilter []\n"
-        "|   |   |   Variable [evalTemp_16]\n"
+        "|   |   |   Variable [evalTemp_80]\n"
         "|   |   PathTraverse []\n"
         "|   |   PathCompare [Eq]\n"
         "|   |   Const [0]\n"
         "|   Filter []\n"
         "|   |   EvalFilter []\n"
-        "|   |   |   Variable [evalTemp_15]\n"
+        "|   |   |   Variable [evalTemp_79]\n"
         "|   |   PathTraverse []\n"
         "|   |   PathCompare [Eq]\n"
         "|   |   Const [0]\n"
         "|   Filter []\n"
         "|   |   EvalFilter []\n"
-        "|   |   |   Variable [evalTemp_14]\n"
+        "|   |   |   Variable [evalTemp_78]\n"
         "|   |   PathTraverse []\n"
         "|   |   PathCompare [Eq]\n"
         "|   |   Const [0]\n"
         "|   Filter []\n"
         "|   |   EvalFilter []\n"
-        "|   |   |   Variable [evalTemp_13]\n"
+        "|   |   |   Variable [evalTemp_77]\n"
         "|   |   PathTraverse []\n"
         "|   |   PathCompare [Eq]\n"
         "|   |   Const [0]\n"
         "|   Filter []\n"
         "|   |   EvalFilter []\n"
-        "|   |   |   Variable [evalTemp_12]\n"
+        "|   |   |   Variable [evalTemp_76]\n"
         "|   |   PathTraverse []\n"
         "|   |   PathCompare [Eq]\n"
         "|   |   Const [0]\n"
@@ -1563,21 +1563,21 @@ TEST(PhysRewriter, FilterIndexingStress) {
         "|   |   limitSkip:\n"
         "|   |       limit: 1\n"
         "|   |       skip: 0\n"
-        "|   Seek [ridProjection: rid_0, {'<root>': root, 'field2': evalTemp_12, 'field5': "
-        "evalTemp_13, 'field6': evalTemp_14, 'field7': evalTemp_15, 'field8': evalTemp_16, "
-        "'field9': evalTemp_17}, c1]\n"
+        "|   Seek [ridProjection: rid_0, {'<root>': root, 'field2': evalTemp_76, 'field5': "
+        "evalTemp_77, 'field6': evalTemp_78, 'field7': evalTemp_79, 'field8': evalTemp_80, "
+        "'field9': evalTemp_81}, c1]\n"
         "|   |   BindBlock:\n"
-        "|   |       [evalTemp_12]\n"
+        "|   |       [evalTemp_76]\n"
         "|   |           Source []\n"
-        "|   |       [evalTemp_13]\n"
+        "|   |       [evalTemp_77]\n"
         "|   |           Source []\n"
-        "|   |       [evalTemp_14]\n"
+        "|   |       [evalTemp_78]\n"
         "|   |           Source []\n"
-        "|   |       [evalTemp_15]\n"
+        "|   |       [evalTemp_79]\n"
         "|   |           Source []\n"
-        "|   |       [evalTemp_16]\n"
+        "|   |       [evalTemp_80]\n"
         "|   |           Source []\n"
-        "|   |       [evalTemp_17]\n"
+        "|   |       [evalTemp_81]\n"
         "|   |           Source []\n"
         "|   |       [root]\n"
         "|   |           Source []\n"
@@ -1585,16 +1585,16 @@ TEST(PhysRewriter, FilterIndexingStress) {
         "|       Variable [rid_0]\n"
         "MergeJoin []\n"
         "|   |   |   Condition\n"
-        "|   |   |       rid_0 = rid_1\n"
+        "|   |   |       rid_0 = rid_2\n"
         "|   |   Collation\n"
         "|   |       Ascending\n"
         "|   Union []\n"
         "|   |   BindBlock:\n"
-        "|   |       [rid_1]\n"
+        "|   |       [rid_2]\n"
         "|   |           Source []\n"
         "|   Evaluation []\n"
         "|   |   BindBlock:\n"
-        "|   |       [rid_1]\n"
+        "|   |       [rid_2]\n"
         "|   |           Variable [rid_0]\n"
         "|   IndexScan [{'<rid>': rid_0}, scanDefName: c1, indexDefName: index3, interval: {[Const "
         "[0], Const [0]], [Const [0], Const [0]]}]\n"
@@ -2253,22 +2253,22 @@ TEST(PhysRewriter, MultiKeyIndex) {
             "|       Variable [rid_0]\n"
             "HashJoin [joinType: Inner]\n"
             "|   |   Condition\n"
-            "|   |       rid_0 = rid_1\n"
+            "|   |       rid_0 = rid_2\n"
             "|   Union []\n"
             "|   |   BindBlock:\n"
-            "|   |       [rid_1]\n"
+            "|   |       [rid_2]\n"
             "|   |           Source []\n"
             "|   Evaluation []\n"
             "|   |   BindBlock:\n"
-            "|   |       [rid_1]\n"
+            "|   |       [rid_2]\n"
             "|   |           Variable [rid_0]\n"
             "|   IndexScan [{'<rid>': rid_0}, scanDefName: c1, indexDefName: index2, interval: "
             "{(Const [2], +inf)}, reversed]\n"
             "|       BindBlock:\n"
             "|           [rid_0]\n"
             "|               Source []\n"
-            "IndexScan [{'<rid>': rid_0}, scanDefName: c1, indexDefName: index1, interval: "
-            "{[Const [1], Const [1]]}]\n"
+            "IndexScan [{'<rid>': rid_0}, scanDefName: c1, indexDefName: index1, interval: {[Const "
+            "[1], Const [1]]}]\n"
             "    BindBlock:\n"
             "        [rid_0]\n"
             "            Source []\n",
@@ -2361,16 +2361,16 @@ TEST(PhysRewriter, CompoundIndex1) {
         "|       Variable [rid_0]\n"
         "MergeJoin []\n"
         "|   |   |   Condition\n"
-        "|   |   |       rid_0 = rid_1\n"
+        "|   |   |       rid_0 = rid_3\n"
         "|   |   Collation\n"
         "|   |       Ascending\n"
         "|   Union []\n"
         "|   |   BindBlock:\n"
-        "|   |       [rid_1]\n"
+        "|   |       [rid_3]\n"
         "|   |           Source []\n"
         "|   Evaluation []\n"
         "|   |   BindBlock:\n"
-        "|   |       [rid_1]\n"
+        "|   |       [rid_3]\n"
         "|   |           Variable [rid_0]\n"
         "|   IndexScan [{'<rid>': rid_0}, scanDefName: c1, indexDefName: index2, interval: {[Const "
         "[2], Const [2]], [Const [4], Const [4]]}]\n"
@@ -2456,6 +2456,7 @@ TEST(PhysRewriter, CompoundIndex2) {
 
     ASSERT_TRUE(phaseManager.optimize(optimized));
     ASSERT_BETWEEN(100, 170, phaseManager.getMemo().getStats()._physPlanExplorationCount);
+
     ASSERT_EXPLAIN_V2(
         "Root []\n"
         "|   |   projections: \n"
@@ -2476,19 +2477,18 @@ TEST(PhysRewriter, CompoundIndex2) {
         "|       Variable [rid_0]\n"
         "MergeJoin []\n"
         "|   |   |   Condition\n"
-        "|   |   |       rid_0 = rid_1\n"
+        "|   |   |       rid_0 = rid_32\n"
         "|   |   Collation\n"
         "|   |       Ascending\n"
         "|   Union []\n"
         "|   |   BindBlock:\n"
-        "|   |       [rid_1]\n"
+        "|   |       [rid_32]\n"
         "|   |           Source []\n"
         "|   Evaluation []\n"
         "|   |   BindBlock:\n"
-        "|   |       [rid_1]\n"
+        "|   |       [rid_32]\n"
         "|   |           Variable [rid_0]\n"
-        "|   IndexScan [{'<rid>': rid_0}, scanDefName: c1, indexDefName: index2, interval: "
-        "{[Const "
+        "|   IndexScan [{'<rid>': rid_0}, scanDefName: c1, indexDefName: index2, interval: {[Const "
         "[2], Const [2]], [Const [4], Const [4]]}]\n"
         "|       BindBlock:\n"
         "|           [rid_0]\n"
@@ -3951,58 +3951,25 @@ TEST(PhysRewriter, IndexPartitioning1) {
     ASSERT_TRUE(phaseManager.optimize(optimized));
     ASSERT_BETWEEN(150, 350, phaseManager.getMemo().getStats()._physPlanExplorationCount);
 
-    ASSERT_EXPLAIN_V2(
-        "Root []\n"
-        "|   |   projections: \n"
-        "|   |       pc\n"
-        "|   RefBlock: \n"
-        "|       Variable [pc]\n"
-        "Exchange []\n"
-        "|   |   distribution: \n"
-        "|   |       type: Centralized\n"
-        "|   RefBlock: \n"
-        "GroupBy []\n"
-        "|   |   groupings: \n"
-        "|   |       RefBlock: \n"
-        "|   |           Variable [pa]\n"
-        "|   aggregations: \n"
-        "|       [pc]\n"
-        "|           Variable [pb]\n"
-        "HashJoin [joinType: Inner]\n"
-        "|   |   Condition\n"
-        "|   |       rid_0 = rid_3\n"
-        "|   Union []\n"
-        "|   |   BindBlock:\n"
-        "|   |       [pa]\n"
-        "|   |           Source []\n"
-        "|   |       [rid_3]\n"
-        "|   |           Source []\n"
-        "|   Evaluation []\n"
-        "|   |   BindBlock:\n"
-        "|   |       [rid_3]\n"
-        "|   |           Variable [rid_0]\n"
-        "|   IndexScan [{'<indexKey> 0': pa, '<rid>': rid_0}, scanDefName: c1, indexDefName: "
-        "index1, interval: {(Const [0], +inf)}]\n"
-        "|       BindBlock:\n"
-        "|           [pa]\n"
-        "|               Source []\n"
-        "|           [rid_0]\n"
-        "|               Source []\n"
-        "Exchange []\n"
-        "|   |   distribution: \n"
-        "|   |       type: HashPartitioning\n"
-        "|   |           projections: \n"
-        "|   |               pa\n"
-        "|   RefBlock: \n"
-        "|       Variable [pa]\n"
-        "IndexScan [{'<indexKey> 0': pb, '<rid>': rid_0}, scanDefName: c1, indexDefName: index2, "
-        "interval: {(Const [1], +inf)}]\n"
-        "    BindBlock:\n"
-        "        [pb]\n"
-        "            Source []\n"
-        "        [rid_0]\n"
-        "            Source []\n",
-        optimized);
+    const BSONObj& result = ExplainGenerator::explainBSONObj(optimized);
+
+    // Compare using BSON since the rid vars are currently unstable for this test.
+    ASSERT_BSON_PATH("\"Root\"", result, "nodeType");
+    ASSERT_BSON_PATH("\"Exchange\"", result, "child.nodeType");
+    ASSERT_BSON_PATH(
+        "{ type: \"Centralized\", disableExchanges: false }", result, "child.distribution");
+    ASSERT_BSON_PATH("\"GroupBy\"", result, "child.child.nodeType");
+    ASSERT_BSON_PATH("\"HashJoin\"", result, "child.child.child.nodeType");
+    ASSERT_BSON_PATH("\"Exchange\"", result, "child.child.child.leftChild.nodeType");
+    ASSERT_BSON_PATH(
+        "{ type: \"HashPartitioning\", disableExchanges: false, projections: [ \"pa\" ] }",
+        result,
+        "child.child.child.leftChild.distribution");
+    ASSERT_BSON_PATH("\"IndexScan\"", result, "child.child.child.leftChild.child.nodeType");
+    ASSERT_BSON_PATH("\"Union\"", result, "child.child.child.rightChild.nodeType");
+    ASSERT_BSON_PATH("\"Evaluation\"", result, "child.child.child.rightChild.children.0.nodeType");
+    ASSERT_BSON_PATH(
+        "\"IndexScan\"", result, "child.child.child.rightChild.children.0.child.nodeType");
 }
 
 TEST(PhysRewriter, LocalGlobalAgg) {
diff --git a/src/mongo/db/query/optimizer/utils/unit_test_utils.h b/src/mongo/db/query/optimizer/utils/unit_test_utils.h
index 402a81618d8..24a32017416 100644
--- a/src/mongo/db/query/optimizer/utils/unit_test_utils.h
+++ b/src/mongo/db/query/optimizer/utils/unit_test_utils.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "mongo/db/bson/dotted_path_support.h"
 #include "mongo/db/operation_context_noop.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
 #include "mongo/db/pipeline/pipeline.h"
@@ -61,6 +62,12 @@ void maybePrintABT(const ABT& abt);
 
 #define ASSERT_EXPLAIN_MEMO(expected, memo) ASSERT_EQ(expected, ExplainGenerator::explainMemo(memo))
 
+#define ASSERT_BSON_PATH(expected, bson, path)                      \
+    ASSERT_EQ(expected,                                             \
+              dotted_path_support::extractElementAtPath(bson, path) \
+                  .toString(false /*includeFieldName*/));
+
+
 #define ASSERT_BETWEEN(a, b, value) \
     ASSERT_LTE(a, value);           \
     ASSERT_GTE(b, value);