SERVER-114845 Get multiplanner into plan ranker (#45019)

GitOrigin-RevId: b5f6ab56796196826c18446337674332ed53a7de
This commit is contained in:
Carlos Alonso 2025-12-11 08:18:07 +01:00 committed by MongoDB Bot
parent ae1e4256dc
commit 8b7abc8214
8 changed files with 101 additions and 60 deletions

View File

@ -97,6 +97,10 @@ WorkingSet* ClassicPlannerInterface::ws() const {
return _plannerData.workingSet.get();
}
std::unique_ptr<WorkingSet> ClassicPlannerInterface::extractWorkingSet() {
return std::move(_plannerData.workingSet);
}
void ClassicPlannerInterface::addDeleteStage(ParsedDelete* parsedDelete,
projection_ast::Projection* projection,
std::unique_ptr<DeleteStageParams> deleteStageParams) {

View File

@ -91,6 +91,11 @@ public:
std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> makeExecutor(
std::unique_ptr<CanonicalQuery> canonicalQuery) final;
/**
* Extracts the WorkingSet used by this planner.
*/
std::unique_ptr<WorkingSet> extractWorkingSet();
protected:
std::unique_ptr<PlanStage> buildExecutableTree(const QuerySolution& qs);

View File

@ -50,11 +50,13 @@ namespace mongo::classic_runtime_planner_for_sbe {
* Data that any runtime planner needs to perform the planning.
*/
struct PlannerDataForSBE final : public PlannerData {
PlannerDataForSBE(OperationContext* opCtx,
PlannerDataForSBE(
OperationContext* opCtx,
CanonicalQuery* cq,
std::unique_ptr<WorkingSet> workingSet,
const MultipleCollectionAccessor& collections,
std::unique_ptr<QueryPlannerParams> plannerParams,
// To be shared between all instances of this type and the prepare helper creating them.
std::shared_ptr<const QueryPlannerParams> plannerParams,
PlanYieldPolicy::YieldPolicy yieldPolicy,
boost::optional<size_t> cachedPlanHash,
std::unique_ptr<PlanYieldPolicySBE> sbeYieldPolicy,

View File

@ -45,7 +45,7 @@ struct PlannerData {
CanonicalQuery* cq,
std::unique_ptr<WorkingSet> workingSet,
const MultipleCollectionAccessor& collections,
std::unique_ptr<QueryPlannerParams> plannerParams,
std::shared_ptr<const QueryPlannerParams> plannerParams,
PlanYieldPolicy::YieldPolicy yieldPolicy,
boost::optional<size_t> cachedPlanHash)
: opCtx(opCtx),
@ -67,7 +67,9 @@ struct PlannerData {
CanonicalQuery* cq;
std::unique_ptr<WorkingSet> workingSet;
const MultipleCollectionAccessor& collections;
std::unique_ptr<QueryPlannerParams> plannerParams;
// Shared pointer since this is shared across all instances of this type and also
// prepare helper functions that indeed create this objects.
std::shared_ptr<const QueryPlannerParams> plannerParams;
PlanYieldPolicy::YieldPolicy yieldPolicy;
boost::optional<size_t> cachedPlanHash;
};

View File

@ -334,20 +334,22 @@ private:
*
* TODO SERVER-87752 Refactor 'PrepareExecutionHelper' to better handle result types.
*/
template <typename KeyType, typename ResultType>
template <typename KeyType, typename ResultType, typename PlannerDataType>
class PrepareExecutionHelper {
public:
PrepareExecutionHelper(OperationContext* opCtx,
const MultipleCollectionAccessor& collections,
PlanYieldPolicy::YieldPolicy yieldPolicy,
CanonicalQuery* cq,
std::unique_ptr<QueryPlannerParams> plannerParams)
std::unique_ptr<QueryPlannerParams> plannerParams,
std::unique_ptr<WorkingSet> workingSet)
: _opCtx{opCtx},
_collections{collections},
_yieldPolicy{yieldPolicy},
_cq{cq},
_plannerParams(std::move(plannerParams)),
_result{std::make_unique<ResultType>()} {
_result{std::make_unique<ResultType>()},
_ws(std::move(workingSet)) {
tassert(11321301, "canonicalQuery must not be null", _cq);
if (shouldLog(MONGO_LOGV2_DEFAULT_COMPONENT, logv2::LogSeverity::Debug(2))) {
_queryStringForDebugLog = _cq->toStringShort();
@ -442,11 +444,12 @@ public:
}
plan_ranking::PlanRanker planRanker;
auto rankerResult =
planRanker.rankPlans(_opCtx, *_cq, *_plannerParams, _yieldPolicy, getCollections());
auto rankerResult = planRanker.rankPlans(
_opCtx, *_cq, *_plannerParams, _yieldPolicy, getCollections(), makePlannerData());
if (!rankerResult.isOK()) {
return rankerResult.getStatus();
}
_ws = planRanker.extractWorkingSet();
std::vector<std::unique_ptr<QuerySolution>> solutions =
std::move(rankerResult.getValue().solutions);
// The planner should have returned an error status if there are no solutions.
@ -575,17 +578,25 @@ protected:
// Stored as a smart pointer for memory safety reasons. Storing a reference would be even
// faster, but also more prone to memory errors. Storing a direct value would incur copying
// costs when std::move()ing, since QueryPlannerParams is just a big aggregated structure.
std::unique_ptr<QueryPlannerParams> _plannerParams;
std::shared_ptr<QueryPlannerParams> _plannerParams;
// In-progress result value of the prepare() call.
std::unique_ptr<ResultType> _result;
// Cached result of CanonicalQuery::toStringShort(). Only populated when logging verbosity is
// high enough to enable messages that need it.
std::string _queryStringForDebugLog;
// WorkingSet for multi planning runs. It is moved into the PlannerData when needed for planning
// like when multiplanning. It is then moved back into this class when planning is done since it
// is needed to build the solution stages.
std::unique_ptr<WorkingSet> _ws;
private:
virtual PlannerDataType makePlannerData() = 0;
};
class ClassicPrepareExecutionHelper final
: public PrepareExecutionHelper<PlanCacheKey, ClassicRuntimePlannerResult> {
: public PrepareExecutionHelper<PlanCacheKey, ClassicRuntimePlannerResult, PlannerData> {
public:
ClassicPrepareExecutionHelper(OperationContext* opCtx,
const MultipleCollectionAccessor& collections,
@ -593,19 +604,19 @@ public:
CanonicalQuery* cq,
PlanYieldPolicy::YieldPolicy yieldPolicy,
std::unique_ptr<QueryPlannerParams> plannerParams)
: PrepareExecutionHelper{opCtx, collections, yieldPolicy, cq, std::move(plannerParams)},
_ws{std::move(ws)} {}
: PrepareExecutionHelper{
opCtx, collections, yieldPolicy, cq, std::move(plannerParams), std::move(ws)} {}
private:
using CachedSolutionPair =
std::pair<std::unique_ptr<CachedSolution>, std::unique_ptr<QuerySolution>>;
PlannerData makePlannerData() {
PlannerData makePlannerData() override {
return PlannerData{_opCtx,
_cq,
std::move(_ws),
_collections,
std::move(_plannerParams),
_plannerParams,
_yieldPolicy,
_cachedPlanHash};
}
@ -729,7 +740,6 @@ private:
return {std::make_pair(std::move(cs), std::move(querySolution))};
}
std::unique_ptr<WorkingSet> _ws;
boost::optional<size_t> _cachedPlanHash;
};
@ -746,7 +756,9 @@ private:
*/
template <class CacheKey, class RuntimePlanningResult>
class SbeWithClassicRuntimePlanningPrepareExecutionHelperBase
: public PrepareExecutionHelper<CacheKey, RuntimePlanningResult> {
: public PrepareExecutionHelper<CacheKey,
RuntimePlanningResult,
classic_runtime_planner_for_sbe::PlannerDataForSBE> {
public:
SbeWithClassicRuntimePlanningPrepareExecutionHelperBase(
OperationContext* opCtx,
@ -754,27 +766,28 @@ public:
std::unique_ptr<WorkingSet> ws,
CanonicalQuery* cq,
PlanYieldPolicy::YieldPolicy policy,
std::unique_ptr<PlanYieldPolicySBE> sbeYieldPolicy,
std::unique_ptr<QueryPlannerParams> plannerParams,
bool useSbePlanCache)
: PrepareExecutionHelper<CacheKey, RuntimePlanningResult>(
opCtx, collections, policy, cq, std::move(plannerParams)),
_ws{std::move(ws)},
_sbeYieldPolicy{std::move(sbeYieldPolicy)},
: PrepareExecutionHelper<CacheKey,
RuntimePlanningResult,
classic_runtime_planner_for_sbe::PlannerDataForSBE>(
opCtx, collections, policy, cq, std::move(plannerParams), std::move(ws)),
_useSbePlanCache{useSbePlanCache} {}
protected:
crp_sbe::PlannerDataForSBE makePlannerData() {
crp_sbe::PlannerDataForSBE makePlannerData() override {
// Use of 'this->' is necessary since some compilers have trouble resolving member
// variables in templated parent class.
return crp_sbe::PlannerDataForSBE{this->_opCtx,
return crp_sbe::PlannerDataForSBE{
this->_opCtx,
this->_cq,
std::move(_ws),
std::move(this->_ws),
this->_collections,
std::move(this->_plannerParams),
this->_plannerParams,
this->_yieldPolicy,
_cachedPlanHash,
std::move(_sbeYieldPolicy),
PlanYieldPolicySBE::make(
this->_opCtx, this->_yieldPolicy, this->_collections, this->_cq->nss()),
_useSbePlanCache};
}
@ -822,13 +835,6 @@ protected:
}
}
std::unique_ptr<WorkingSet> _ws;
// When using the classic multi-planner for SBE, we need both classic and SBE yield policy to
// support yielding during trial period in classic engine. The classic yield policy to stored in
// 'PrepareExecutionHelper'.
std::unique_ptr<PlanYieldPolicySBE> _sbeYieldPolicy;
const bool _useSbePlanCache;
// If there is a matching cache entry, this is the hash of that plan.
@ -849,14 +855,12 @@ public:
std::unique_ptr<WorkingSet> ws,
CanonicalQuery* cq,
PlanYieldPolicy::YieldPolicy policy,
std::unique_ptr<PlanYieldPolicySBE> sbeYieldPolicy,
std::unique_ptr<QueryPlannerParams> plannerParams)
: SbeWithClassicRuntimePlanningPrepareExecutionHelperBase{opCtx,
collections,
std::move(ws),
cq,
policy,
std::move(sbeYieldPolicy),
std::move(plannerParams),
true /*useSbePlanCache*/} {}
@ -928,14 +932,12 @@ public:
std::unique_ptr<WorkingSet> ws,
CanonicalQuery* cq,
PlanYieldPolicy::YieldPolicy policy,
std::unique_ptr<PlanYieldPolicySBE> sbeYieldPolicy,
std::unique_ptr<QueryPlannerParams> plannerParams)
: SbeWithClassicRuntimePlanningPrepareExecutionHelperBase{opCtx,
collections,
std::move(ws),
cq,
policy,
std::move(sbeYieldPolicy),
std::move(plannerParams),
false /*useSbePlanCache*/} {}
@ -1029,7 +1031,6 @@ std::unique_ptr<PlannerInterface> getClassicPlannerForSbe(
const MultipleCollectionAccessor& collections,
CanonicalQuery* canonicalQuery,
PlanYieldPolicy::YieldPolicy yieldPolicy,
std::unique_ptr<PlanYieldPolicySBE> sbeYieldPolicy,
std::unique_ptr<QueryPlannerParams> plannerParams) {
PrepareExecutionHelperType helper{
opCtx,
@ -1037,7 +1038,6 @@ std::unique_ptr<PlannerInterface> getClassicPlannerForSbe(
std::make_unique<WorkingSet>(),
std::move(canonicalQuery),
yieldPolicy,
std::move(sbeYieldPolicy),
std::move(plannerParams),
};
ScopedDebugInfo queryPlannerParams(
@ -1366,9 +1366,6 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind
// If we have a distinct, we might get a better plan using classic and DISTINCT_SCAN than
// SBE without one.
if (useSbeEngine && canCommitToSbe()) {
auto sbeYieldPolicy =
PlanYieldPolicySBE::make(opCtx, yieldPolicy, collections, canonicalQuery->nss());
plannerParams->fillOutSecondaryCollectionsPlannerParams(
opCtx, *canonicalQuery, collections);
@ -1382,7 +1379,6 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind
collections,
canonicalQuery.get(),
yieldPolicy,
std::move(sbeYieldPolicy),
std::move(plannerParams));
} else {
canonicalQuery->setUsingSbePlanCache(false);
@ -1392,7 +1388,6 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind
collections,
canonicalQuery.get(),
yieldPolicy,
std::move(sbeYieldPolicy),
std::move(plannerParams));
}
}

View File

@ -45,12 +45,13 @@
#include <algorithm>
#include <cstddef>
#include <deque>
#include <functional>
#include <memory>
#include <string>
#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery
// TODO SERVER-115240. Rename and maybe move it to a separate plan scoring directory.
namespace mongo::plan_ranker {
// Constant used for tie breakers.

View File

@ -30,6 +30,7 @@
#include "mongo/db/query/plan_ranking/plan_ranker.h"
#include "mongo/base/status.h"
#include "mongo/db/exec/runtime_planners/classic_runtime_planner/planner_interface.h"
#include "mongo/db/query/canonical_query.h"
#include "mongo/db/query/multiple_collection_accessor.h"
#include "mongo/db/query/plan_ranking/cbr_plan_ranking.h"
@ -44,9 +45,23 @@ StatusWith<QueryPlanner::PlanRankingResult> PlanRanker::rankPlans(
CanonicalQuery& query,
const QueryPlannerParams& plannerParams,
PlanYieldPolicy::YieldPolicy yieldPolicy,
const MultipleCollectionAccessor& collections) const {
const MultipleCollectionAccessor& collections,
PlannerData plannerData) {
// TODO SERVER-114514: Implement restrictive approach for automaticCE. Then delete this comment
// below:
// For illustration purposes. We'd use multiplanner for ranking like this:
// auto statusWithMultiPlanSolns = QueryPlanner::plan(query, plannerParams);
// auto mp =
// classic_runtime_planner::MultiPlanner(std::move(plannerData),
// std::move(statusWithMultiPlanSolns.getValue()),
// QueryPlanner::PlanRankingResult{});
// uassertStatusOK(mp.plan()); <- We could also use runTrials or any other method available.
// Once done, extract the WorkingSet from the MultiPlanner and store it here for further
// extraction.
// _ws.emplace(mp.extractWorkingSet());
auto rankerMode = plannerParams.planRankerMode;
if (rankerMode != QueryPlanRankerModeEnum::kMultiPlanning) {
_ws = std::move(plannerData.workingSet);
return CBRPlanRankingStrategy().rankPlans(
opCtx, query, plannerParams, yieldPolicy, collections);
} else {
@ -56,6 +71,7 @@ StatusWith<QueryPlanner::PlanRankingResult> PlanRanker::rankPlans(
* to select a winning plan at runtime.
*/
auto statusWithMultiPlanSolns = QueryPlanner::plan(query, plannerParams);
_ws = std::move(plannerData.workingSet);
if (!statusWithMultiPlanSolns.isOK()) {
return statusWithMultiPlanSolns.getStatus().withContext(
str::stream() << "error processing query: " << query.toStringForErrorMsg()
@ -64,5 +80,12 @@ StatusWith<QueryPlanner::PlanRankingResult> PlanRanker::rankPlans(
return QueryPlanner::PlanRankingResult{std::move(statusWithMultiPlanSolns.getValue())};
}
}
std::unique_ptr<WorkingSet> PlanRanker::extractWorkingSet() {
tassert(11484500, "WorkingSet is not initialized", _ws);
auto result = std::move(_ws);
_ws = nullptr;
return result;
}
} // namespace plan_ranking
} // namespace mongo

View File

@ -29,6 +29,7 @@
#pragma once
#include "mongo/db/exec/runtime_planners/planner_interface.h"
#include "mongo/db/query/canonical_query.h"
#include "mongo/db/query/multiple_collection_accessor.h"
#include "mongo/db/query/plan_yield_policy.h"
@ -53,7 +54,15 @@ public:
CanonicalQuery& query,
const QueryPlannerParams& plannerParams,
PlanYieldPolicy::YieldPolicy yieldPolicy,
const MultipleCollectionAccessor& collections) const;
const MultipleCollectionAccessor& collections,
// PlannerData for classic multiplanner. We only need the classic one since multiplanning
// only runs with classic, even if SBE is enabled.
PlannerData multiPlannerData);
std::unique_ptr<WorkingSet> extractWorkingSet();
private:
std::unique_ptr<WorkingSet> _ws;
};
} // namespace plan_ranking
} // namespace mongo