SERVER-54083 Include scores calculated by PlanRanker in explain "queryPlanner" output and/or "allPlansExecution"

6 files changed, 72 insertions, 19 deletions

diff --git a/src/mongo/db/query/explain.cpp b/src/mongo/db/query/explain.cpp
index f40c7a6b1c0..f2052991556 100644
--- a/src/mongo/db/query/explain.cpp
+++ b/src/mongo/db/query/explain.cpp
@@ -159,11 +159,16 @@ void generatePlannerInfo(PlanExecutor* exec,
  * section, but will not affect the reporting of timing for individual stages. If 'totalTimeMillis'
  * is not set, we use the approximate timing information collected by the stages.
  *
+ * The 'isTrialPeriodInfo' value indicates whether the function was called to generate the
+ * stats collected during the trial period of the plan selection phase, i.e is this section being
+ * generated for the 'allPlansExecution' field.
+ *
  * Stats are generated at the verbosity specified by 'verbosity'.
  */
 void generateSinglePlanExecutionInfo(const PlanExplainer::PlanStatsDetails& details,
                                      boost::optional<long long> totalTimeMillis,
-                                     BSONObjBuilder* out) {
+                                     BSONObjBuilder* out,
+                                     bool isTrialPeriodInfo) {
     auto&& [stats, summary] = details;
     invariant(summary);
 
@@ -183,6 +188,12 @@ void generateSinglePlanExecutionInfo(const PlanExplainer::PlanStatsDetails& deta
         out->appendBool("failed", true);
     }
 
+    // Only the scores calculated from the trial period should be outputted alongside each plan
+    // in 'allPlansExecution' and not alongside the winning plan stats in 'executionStats'.
+    if (isTrialPeriodInfo && summary->score) {
+        out->appendNumber("score", *summary->score);
+    }
+
     // Add the tree of stages, with individual execution stats for each stage.
     out->append("executionStages", stats);
 }
@@ -222,8 +233,10 @@ void generateExecutionInfo(PlanExecutor* exec,
     // Generate exec stats BSON for the winning plan.
     auto opCtx = exec->getOpCtx();
     auto totalTimeMillis = durationCount<Milliseconds>(CurOp::get(opCtx)->elapsedTimeTotal());
-    generateSinglePlanExecutionInfo(
-        explainer.getWinningPlanStats(verbosity), totalTimeMillis, &execBob);
+    generateSinglePlanExecutionInfo(explainer.getWinningPlanStats(verbosity),
+                                    totalTimeMillis,
+                                    &execBob,
+                                    false /* isTrialPeriodInfo */);
 
     // Also generate exec stats for all plans, if the verbosity level is high enough. These stats
     // reflect what happened during the trial period that ranked the plans.
@@ -237,12 +250,14 @@ void generateExecutionInfo(PlanExecutor* exec,
         // If the winning plan was uncontested, leave the `allPlansExecution` array empty.
         if (explainer.isMultiPlan()) {
             BSONObjBuilder planBob(allPlansBob.subobjStart());
-            generateSinglePlanExecutionInfo(*winningPlanTrialStats, boost::none, &planBob);
+            generateSinglePlanExecutionInfo(
+                *winningPlanTrialStats, boost::none, &planBob, true /* isTrialPeriodInfo */);
             planBob.doneFast();
 
             for (auto&& stats : explainer.getRejectedPlansStats(verbosity)) {
                 BSONObjBuilder planBob(allPlansBob.subobjStart());
-                generateSinglePlanExecutionInfo(stats, boost::none, &planBob);
+                generateSinglePlanExecutionInfo(
+                    stats, boost::none, &planBob, true /* isTrialPeriodInfo */);
                 planBob.doneFast();
             }
         }
@@ -414,7 +429,8 @@ void Explain::planCacheEntryToBSON(const PlanCacheEntry& entry, BSONObjBuilder*
         BSONArrayBuilder creationBuilder(out->subarrayStart("creationExecStats"));
         for (auto&& stats : execStats) {
             BSONObjBuilder planBob(creationBuilder.subobjStart());
-            generateSinglePlanExecutionInfo(stats, boost::none, &planBob);
+            generateSinglePlanExecutionInfo(
+                stats, boost::none, &planBob, false /* isTrialPeriodInfo */);
             planBob.doneFast();
         }
         creationBuilder.doneFast();
diff --git a/src/mongo/db/query/plan_explainer_impl.cpp b/src/mongo/db/query/plan_explainer_impl.cpp
index e832698de78..9720f628292 100644
--- a/src/mongo/db/query/plan_explainer_impl.cpp
+++ b/src/mongo/db/query/plan_explainer_impl.cpp
@@ -639,6 +639,20 @@ const boost::optional<size_t> getWinningPlanIdx(PlanStage* root) {
     return {};
 }
 
+/**
+ * If 'root' has a MultiPlanStage returns the score of its best plan.
+ */
+boost::optional<double> getWinningPlanScore(PlanStage* root) {
+    if (const auto mps = getMultiPlanStage(root); mps) {
+        auto bestPlanIdx = mps->bestPlanIdx();
+        tassert(5408300,
+                "Trying to get best plan index of a MultiPlanStage without winning plan",
+                bestPlanIdx);
+        return mps->getCandidateScore(*bestPlanIdx);
+    }
+    return {};
+}
+
 void PlanExplainerImpl::getSummaryStats(PlanSummaryStats* statsOut) const {
     invariant(statsOut);
 
@@ -725,7 +739,11 @@ PlanExplainer::PlanStatsDetails PlanExplainerImpl::getWinningPlanStats(
         -> std::pair<std::unique_ptr<PlanStageStats>, const boost::optional<PlanSummaryStats>> {
         auto stats = _root->getStats();
         if (verbosity >= ExplainOptions::Verbosity::kExecStats) {
-            return {std::move(stats), collectExecutionStatsSummary(stats.get(), winningPlanIdx)};
+            auto summary = collectExecutionStatsSummary(stats.get(), winningPlanIdx);
+            if (verbosity >= ExplainOptions::Verbosity::kExecAllPlans) {
+                summary.score = getWinningPlanScore(_root);
+            }
+            return {std::move(stats), summary};
         }
 
         return {std::move(stats), boost::none};
@@ -737,7 +755,7 @@ PlanExplainer::PlanStatsDetails PlanExplainerImpl::getWinningPlanStats(
 }
 
 PlanExplainer::PlanStatsDetails PlanExplainerImpl::getWinningPlanTrialStats() const {
-    return getWinningPlanStats(ExplainOptions::Verbosity::kExecStats);
+    return getWinningPlanStats(ExplainOptions::Verbosity::kExecAllPlans);
 }
 
 std::vector<PlanExplainer::PlanStatsDetails> PlanExplainerImpl::getRejectedPlansStats(
@@ -758,9 +776,17 @@ std::vector<PlanExplainer::PlanStatsDetails> PlanExplainerImpl::getRejectedPlans
             BSONObjBuilder bob;
             auto stats = _root->getStats();
             statsToBSON(*stats, verbosity, i, &bob, &bob);
-            res.push_back({bob.obj(),
-                           {verbosity >= ExplainOptions::Verbosity::kExecStats,
-                            collectExecutionStatsSummary(stats.get(), i)}});
+            auto summary = [&]() -> boost::optional<PlanSummaryStats> {
+                if (verbosity >= ExplainOptions::Verbosity::kExecStats) {
+                    auto summary = collectExecutionStatsSummary(stats.get(), i);
+                    if (verbosity >= ExplainOptions::Verbosity::kExecAllPlans) {
+                        summary.score = mps->getCandidateScore(i);
+                    }
+                    return summary;
+                }
+                return {};
+            }();
+            res.push_back({bob.obj(), summary});
         }
     }
 
diff --git a/src/mongo/db/query/plan_explainer_sbe.cpp b/src/mongo/db/query/plan_explainer_sbe.cpp
index 63f9a9fff86..67690127987 100644
--- a/src/mongo/db/query/plan_explainer_sbe.cpp
+++ b/src/mongo/db/query/plan_explainer_sbe.cpp
@@ -299,7 +299,7 @@ PlanSummaryStats collectExecutionStatsSummary(const sbe::PlanStageStats* stats)
 }
 
 PlanExplainer::PlanStatsDetails buildPlanStatsDetails(
-    const QuerySolutionNode* node,
+    const QuerySolution* solution,
     const sbe::PlanStageStats* stats,
     const boost::optional<BSONObj>& execPlanDebugInfo,
     ExplainOptions::Verbosity verbosity) {
@@ -307,6 +307,9 @@ PlanExplainer::PlanStatsDetails buildPlanStatsDetails(
 
     if (verbosity >= ExplainOptions::Verbosity::kExecStats) {
         auto summary = collectExecutionStatsSummary(stats);
+        if (verbosity >= ExplainOptions::Verbosity::kExecAllPlans) {
+            summary.score = solution->score;
+        }
         statsToBSON(stats, &bob, &bob);
         // At the 'kQueryPlanner' verbosity level we use the QSN-derived format for the given plan,
         // and thus the winning plan and rejected plans at this verbosity should display the
@@ -317,7 +320,7 @@ PlanExplainer::PlanStatsDetails buildPlanStatsDetails(
         return {bob.obj(), std::move(summary)};
     }
 
-    statsToBSON(node, &bob, &bob);
+    statsToBSON(solution->root(), &bob, &bob);
     invariant(execPlanDebugInfo);
     return {BSON("queryPlan" << bob.obj() << "slotBasedPlan" << *execPlanDebugInfo), boost::none};
 }
@@ -484,7 +487,7 @@ PlanExplainer::PlanStatsDetails PlanExplainerSBE::getWinningPlanStats(
     invariant(_solution);
     auto stats = _root->getStats(true /* includeDebugInfo  */);
     return buildPlanStatsDetails(
-        _solution->root(), stats.get(), buildExecPlanDebugInfo(_root, _rootData), verbosity);
+        _solution, stats.get(), buildExecPlanDebugInfo(_root, _rootData), verbosity);
 }
 
 PlanExplainer::PlanStatsDetails PlanExplainerSBE::getWinningPlanTrialStats() const {
@@ -492,14 +495,14 @@ PlanExplainer::PlanStatsDetails PlanExplainerSBE::getWinningPlanTrialStats() con
     if (_rootData->savedStatsOnEarlyExit) {
         invariant(_solution);
         return buildPlanStatsDetails(
-            _solution->root(),
+            _solution,
             _rootData->savedStatsOnEarlyExit.get(),
             // This parameter is not used in `buildPlanStatsDetails` if the last parameter is
-            // `ExplainOptions::Verbosity::kExecStats`, as is the case here.
+            // `ExplainOptions::Verbosity::kExecAllPlans`, as is the case here.
             boost::none,
-            ExplainOptions::Verbosity::kExecStats);
+            ExplainOptions::Verbosity::kExecAllPlans);
     }
-    return getWinningPlanStats(ExplainOptions::Verbosity::kExecStats);
+    return getWinningPlanStats(ExplainOptions::Verbosity::kExecAllPlans);
 }
 
 std::vector<PlanExplainer::PlanStatsDetails> PlanExplainerSBE::getRejectedPlansStats(
@@ -517,7 +520,7 @@ std::vector<PlanExplainer::PlanStatsDetails> PlanExplainerSBE::getRejectedPlansS
         auto stats = candidate.root->getStats(true /* includeDebugInfo  */);
         auto execPlanDebugInfo = buildExecPlanDebugInfo(candidate.root.get(), &candidate.data);
         res.push_back(buildPlanStatsDetails(
-            candidate.solution->root(), stats.get(), execPlanDebugInfo, verbosity));
+            candidate.solution.get(), stats.get(), execPlanDebugInfo, verbosity));
     }
     return res;
 }
diff --git a/src/mongo/db/query/plan_ranker_util.h b/src/mongo/db/query/plan_ranker_util.h
index 0916fa703af..4874969db0d 100644
--- a/src/mongo/db/query/plan_ranker_util.h
+++ b/src/mongo/db/query/plan_ranker_util.h
@@ -115,6 +115,7 @@ StatusWith<std::unique_ptr<PlanRankingDecision>> pickBestPlan(
                 score += 1;
             }
 
+            candidates[i].solution->score = score;
             scoresAndCandidateIndices.push_back(std::make_pair(score, i));
         } else {
             failed.push_back(i);
diff --git a/src/mongo/db/query/plan_summary_stats.h b/src/mongo/db/query/plan_summary_stats.h
index cfaa2053d16..129c4f83895 100644
--- a/src/mongo/db/query/plan_summary_stats.h
+++ b/src/mongo/db/query/plan_summary_stats.h
@@ -112,6 +112,10 @@ struct PlanSummaryStats {
 
     // Was a replan triggered during the execution of this query?
     std::optional<std::string> replanReason;
+
+    // Score calculated for the plan by PlanRanker. Only set if there were multiple candidate plans
+    // and allPlansExecution verbosity mode is selected.
+    boost::optional<double> score;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index c0d3411dd44..85c41b2b598 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -382,6 +382,9 @@ public:
 
     PlanEnumeratorExplainInfo _enumeratorExplainInfo;
 
+    // Score calculated by PlanRanker. Only present if there are multiple candidate plans.
+    boost::optional<double> score;
+
 private:
     using QsnIdGenerator = IdGenerator<PlanNodeId>;