SERVER-99617 Boost DISTINCT_SCAN productivity score (#31618)

alyacb · MongoDB Bot · commit 7379d610923c · 2025-02-05T14:29:19.000Z
GitOrigin-RevId: 5d7295b
diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
@@ -267,7 +267,7 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {
 
     // After picking best plan, ranking will own plan stats from candidate solutions (winner and
     // losers).
-    auto statusWithRanking = plan_ranker::pickBestPlan(_candidates);
+    auto statusWithRanking = plan_ranker::pickBestPlan(_candidates, *_query);
     if (!statusWithRanking.isOK()) {
         return statusWithRanking.getStatus();
     }
diff --git a/src/mongo/db/query/plan_ranker.cpp b/src/mongo/db/query/plan_ranker.cpp
@@ -57,7 +57,8 @@ void logScoreFormula(std::function<std::string()> formula,
                      double noFetchBonus,
                      double noSortBonus,
                      double noIxisectBonus,
-                     double tieBreakers) {
+                     double tieBreakers,
+                     boost::optional<double> groupByDistinctBonus) {
     LOGV2_DEBUG(
         20961, 2, "Score formula", "formula"_attr = [&]() {
             StringBuilder sb;
@@ -69,6 +70,9 @@ void logScoreFormula(std::function<std::string()> formula,
                << " noFetchBonus + " << str::convertDoubleToString(noSortBonus) << " noSortBonus + "
                << str::convertDoubleToString(noIxisectBonus)
                << " noIxisectBonus = " << str::convertDoubleToString(tieBreakers) << ")";
+            if (groupByDistinctBonus) {
+                sb << " + groupByDistinctBonus(" << *groupByDistinctBonus << ")";
+            }
             return sb.str();
         }());
 }
@@ -77,6 +81,14 @@ void logScoreBoost(double score) {
     LOGV2_DEBUG(20962, 5, "Score boosted due to intersection forcing", "newScore"_attr = score);
 }
 
+void logScoreGroupByDistinctBoost(double bonus) {
+    LOGV2_DEBUG(
+        9961700,
+        5,
+        "Adding groupByDistinctBonus, boost formula is: std::min(1 - productivity, productivity)",
+        "groupByDistinctBonus"_attr = bonus);
+}
+
 void logScoringPlan(std::function<std::string()> solution,
                     std::function<std::string()> explain,
                     std::function<std::string()> planSummary,
diff --git a/src/mongo/db/query/plan_ranker.h b/src/mongo/db/query/plan_ranker.h
@@ -69,8 +69,10 @@ void logScoreFormula(std::function<std::string()> formula,
                      double noFetchBonus,
                      double noSortBonus,
                      double noIxisectBonus,
-                     double tieBreakers);
+                     double tieBreakers,
+                     boost::optional<double> groupByDistinctBonus);
 void logScoreBoost(double score);
+void logScoreGroupByDistinctBoost(double bonus);
 void logScoringPlan(std::function<std::string()> solution,
                     std::function<std::string()> explain,
                     std::function<std::string()> planSummary,
@@ -103,7 +105,7 @@ class PlanScorer {
     PlanScorer() = default;
     virtual ~PlanScorer() = default;
 
-    double calculateScore(const PlanStageStatsType* stats) const {
+    double calculateScore(const PlanStageStatsType* stats, const CanonicalQuery& cq) const {
         // We start all scores at 1.  Our "no plan selected" score is 0 and we want all plans to
         // be greater than that.
         const double baseScore = 1;
@@ -139,7 +141,23 @@ class PlanScorer {
         }
 
         const double tieBreakers = noFetchBonus + noSortBonus + noIxisectBonus;
-        double score = baseScore + productivity + tieBreakers;
+        boost::optional<double> groupByDistinctBonus;
+
+        // Apply a large bonus to DISTINCT_SCAN plans in an aggregaton context, as the
+        // $groupByDistinct rewrite can reduce the amount of overall work the query needs to do.
+        if (cq.getExpCtx()->isFeatureFlagShardFilteringDistinctScanEnabled() && cq.getDistinct() &&
+            !cq.cqPipeline().empty() && hasStage(STAGE_DISTINCT_SCAN, stats)) {
+            // Assume that every advance in a distinct scan is twice as productive as the
+            // equivalent index scan, up to the number of works actually done by the
+            // distinct scan, in order to favor distinct scans. The maximum bonus is 0.5
+            // (productivity = 0.5), while the minimum bonus is 0 (productivity = 1). If the
+            // distinct scan is not very productive (< 0.5) we don't want to prioritize it
+            // too much; conversely, if it is very productive, we don't need a huge bonus.
+            groupByDistinctBonus = std::min(1 - productivity, productivity);
+            log_detail::logScoreGroupByDistinctBoost(*groupByDistinctBonus);
+        }
+
+        double score = baseScore + productivity + tieBreakers + groupByDistinctBonus.value_or(0.0);
 
         log_detail::logScoreFormula([this, stats] { return getProductivityFormula(stats); },
                                     score,
@@ -148,7 +166,8 @@ class PlanScorer {
                                     noFetchBonus,
                                     noSortBonus,
                                     noIxisectBonus,
-                                    tieBreakers);
+                                    tieBreakers,
+                                    groupByDistinctBonus);
 
         if (internalQueryForceIntersectionPlans.load()) {
             if (hasStage(STAGE_AND_HASH, stats) || hasStage(STAGE_AND_SORTED, stats)) {
diff --git a/src/mongo/db/query/plan_ranker_test.cpp b/src/mongo/db/query/plan_ranker_test.cpp
@@ -31,13 +31,19 @@
  * This file contains tests for mongo/db/query/plan_ranker.h
  */
 
-#include "mongo/db/query/plan_ranker.h"
-
 #include <utility>
 #include <vector>
 
+#include "mongo/bson/bsonelement.h"
+#include "mongo/bson/bsonmisc.h"
 #include "mongo/db/exec/plan_stats.h"
+#include "mongo/db/namespace_string.h"
+#include "mongo/db/pipeline/document_source_group.h"
+#include "mongo/db/pipeline/expression_context_for_test.h"
+#include "mongo/db/query/canonical_query.h"
 #include "mongo/db/query/plan_ranker_util.h"
+#include "mongo/db/query/stage_types.h"
+#include "mongo/idl/server_parameter_test_util.h"
 #include "mongo/unittest/assert.h"
 #include "mongo/unittest/framework.h"
 
@@ -49,11 +55,21 @@ using std::make_unique;
 using std::string;
 using std::unique_ptr;
 
+unique_ptr<CanonicalQuery> makeCanonicalQuery() {
+    auto expCtx = new ExpressionContextForTest();
+    auto findCommand = std::make_unique<FindCommandRequest>(NamespaceString());
+    return std::make_unique<CanonicalQuery>(CanonicalQueryParams{
+        .expCtx = expCtx, .parsedFind = ParsedFindCommandParams{std::move(findCommand)}});
+}
+
 unique_ptr<PlanStageStats> makeStats(const char* name,
                                      StageType type,
-                                     unique_ptr<SpecificStats> specific) {
+                                     unique_ptr<SpecificStats> specific,
+                                     size_t works = 1,
+                                     size_t advances = 1) {
     auto stats = make_unique<PlanStageStats>(name, type);
-    stats->common.works = 1;
+    stats->common.works = works;
+    stats->common.advanced = advances;
     stats->specific = std::move(specific);
     return stats;
 }
@@ -74,11 +90,59 @@ TEST(PlanRankerTest, NoFetchBonus) {
     badPlan->children[0]->children.emplace_back(
         makeStats("IXSCAN", STAGE_IXSCAN, make_unique<IndexScanStats>()));
 
+    auto cq = makeCanonicalQuery();
     auto scorer = plan_ranker::makePlanScorer();
-    auto goodScore = scorer->calculateScore(goodPlan.get());
-    auto badScore = scorer->calculateScore(badPlan.get());
+    auto goodScore = scorer->calculateScore(goodPlan.get(), *cq);
+    auto badScore = scorer->calculateScore(badPlan.get(), *cq);
 
     ASSERT_GT(goodScore, badScore);
 }
 
+TEST(PlanRankerTest, DistinctBonus) {
+    RAIIServerParameterControllerForTest shardFilteringDistinct(
+        "featureFlagShardFilteringDistinctScan", true);
+
+    // Two plans: both fetch, one is a DISTINCT_SCAN, other is an IXSCAN.
+    // DISTINCT_SCAN does 2 advances / 10 works.
+    auto dsStats = make_unique<DistinctScanStats>();
+    dsStats->isFetching = true;
+    dsStats->isShardFilteringDistinctScanEnabled = true;
+    auto distinctScanPlan =
+        makeStats("DISTINCT_SCAN", STAGE_DISTINCT_SCAN, std::move(dsStats), 10, 2);
+
+    // IXSCAN plan does 2 advances / 10 works.
+    auto ixscanPlan = makeStats("FETCH", STAGE_FETCH, make_unique<FetchStats>(), 10, 2);
+    ixscanPlan->children.emplace_back(
+        makeStats("IXSCAN", STAGE_IXSCAN, make_unique<IndexScanStats>(), 10, 2));
+
+    auto cq = makeCanonicalQuery();
+    cq->setDistinct(CanonicalDistinct("someKey"));
+    auto scorer = plan_ranker::makePlanScorer();
+    auto distinctScore = scorer->calculateScore(distinctScanPlan.get(), *cq);
+    auto ixscanScore = scorer->calculateScore(ixscanPlan.get(), *cq);
+
+    // Both plans should tie now- a tie-breaker will be applied at a later stage.
+    ASSERT_EQ(distinctScore, ixscanScore);
+
+    // Now we change to an aggregation context (simulate $groupByDistinct rewrite case).
+    auto groupBson = BSON("$group" << BSON("_id"
+                                           << "someKey"));
+    cq->setCqPipeline(
+        {DocumentSourceGroup::createFromBson(groupBson.firstElement(), cq->getExpCtx())}, true);
+
+    // When in a distinct() context, productivity is considered larger in a distinct, even if both
+    // plans have the same advances:work ratio. A DISTINCT_SCAN should now win by a large margin
+    // (tie breaker).
+    distinctScore = scorer->calculateScore(distinctScanPlan.get(), *cq);
+    ixscanScore = scorer->calculateScore(ixscanPlan.get(), *cq);
+    ASSERT_GT(distinctScore, ixscanScore);
+
+    // If we make the IXSCAN productive enough, however, it can still win!
+    ixscanPlan->children[0]->common.advanced = 10;
+    ixscanPlan->common.advanced = 10;
+    distinctScore = scorer->calculateScore(distinctScanPlan.get(), *cq);
+    ixscanScore = scorer->calculateScore(ixscanPlan.get(), *cq);
+    ASSERT_GT(ixscanScore, distinctScore);
+}
+
 };  // namespace
diff --git a/src/mongo/db/query/plan_ranker_util.h b/src/mongo/db/query/plan_ranker_util.h
@@ -208,7 +208,8 @@ void addTieBreakingHeuristicsBonuses(
  */
 template <typename PlanStageType, typename ResultType, typename Data>
 StatusWith<std::unique_ptr<PlanRankingDecision>> pickBestPlan(
-    const std::vector<BaseCandidatePlan<PlanStageType, ResultType, Data>>& candidates) {
+    const std::vector<BaseCandidatePlan<PlanStageType, ResultType, Data>>& candidates,
+    const CanonicalQuery& cq) {
     invariant(!candidates.empty());
     // A plan that hits EOF is automatically scored above
     // its peers. If multiple plans hit EOF during the same
@@ -243,7 +244,7 @@ StatusWith<std::unique_ptr<PlanRankingDecision>> pickBestPlan(
                                        [&]() { return explainer->getPlanSummary(); },
                                        i,
                                        statTrees[i]->common.isEOF);
-            double score = makePlanScorer()->calculateScore(statTrees[i].get());
+            double score = makePlanScorer()->calculateScore(statTrees[i].get(), cq);
             log_detail::logScore(score);
             if (statTrees[i]->common.isEOF) {
                 log_detail::logEOFBonus(eofBonus);

Original file line number	Diff line number	Diff line change
`@@ -267,7 +267,7 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {`
`267`	`267`
`268`	`268`	`// After picking best plan, ranking will own plan stats from candidate solutions (winner and`
`269`	`269`	`// losers).`
`270`		`- auto statusWithRanking = plan_ranker::pickBestPlan(_candidates);`
	`270`	`+ auto statusWithRanking = plan_ranker::pickBestPlan(_candidates, *_query);`
`271`	`271`	`if (!statusWithRanking.isOK()) {`
`272`	`272`	`return statusWithRanking.getStatus();`
`273`	`273`	`}`