diff options
Diffstat (limited to 'src/mongo/db/exec')
76 files changed, 13698 insertions, 618 deletions
diff --git a/src/mongo/db/exec/SConscript b/src/mongo/db/exec/SConscript index d0d64ecc9d7..0114f29580f 100644 --- a/src/mongo/db/exec/SConscript +++ b/src/mongo/db/exec/SConscript @@ -7,6 +7,7 @@ env = env.Clone() env.SConscript( dirs=[ "document_value", + "sbe", ], exports=[ "env", diff --git a/src/mongo/db/exec/cached_plan.cpp b/src/mongo/db/exec/cached_plan.cpp index 495ba879da4..3933f1aa938 100644 --- a/src/mongo/db/exec/cached_plan.cpp +++ b/src/mongo/db/exec/cached_plan.cpp @@ -38,7 +38,9 @@ #include "mongo/db/catalog/collection.h" #include "mongo/db/concurrency/write_conflict_exception.h" #include "mongo/db/exec/multi_plan.h" +#include "mongo/db/exec/plan_cache_util.h" #include "mongo/db/exec/scoped_timer.h" +#include "mongo/db/exec/trial_period_utils.h" #include "mongo/db/exec/working_set_common.h" #include "mongo/db/query/collection_query_info.h" #include "mongo/db/query/explain.h" @@ -46,7 +48,7 @@ #include "mongo/db/query/plan_yield_policy.h" #include "mongo/db/query/query_knobs_gen.h" #include "mongo/db/query/query_planner.h" -#include "mongo/db/query/stage_builder.h" +#include "mongo/db/query/stage_builder_util.h" #include "mongo/logv2/log.h" #include "mongo/util/str.h" #include "mongo/util/transitional_tools_do_not_use/vector_spooling.h" @@ -91,7 +93,7 @@ Status CachedPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { static_cast<size_t>(internalQueryCacheEvictionRatio * _decisionWorks); // The trial period ends without replanning if the cached plan produces this many results. - size_t numResults = MultiPlanStage::getTrialPeriodNumToReturn(*_canonicalQuery); + size_t numResults = trial_period::getTrialPeriodNumToReturn(*_canonicalQuery); for (size_t i = 0; i < maxWorksBeforeReplan; ++i) { // Might need to yield between calls to work due to the timer elapsing. @@ -186,9 +188,9 @@ Status CachedPlanStage::tryYield(PlanYieldPolicy* yieldPolicy) { // 2) some stage requested a yield, or // 3) we need to yield and retry due to a WriteConflictException. // In all cases, the actual yielding happens here. - if (yieldPolicy->shouldYieldOrInterrupt()) { + if (yieldPolicy->shouldYieldOrInterrupt(expCtx()->opCtx)) { // Here's where we yield. - return yieldPolicy->yieldOrInterrupt(); + return yieldPolicy->yieldOrInterrupt(expCtx()->opCtx); } return Status::OK(); @@ -222,8 +224,8 @@ Status CachedPlanStage::replan(PlanYieldPolicy* yieldPolicy, bool shouldCache, s if (1 == solutions.size()) { // Only one possible plan. Build the stages from the solution. - auto newRoot = - StageBuilder::build(opCtx(), collection(), *_canonicalQuery, *solutions[0], _ws); + auto&& newRoot = stage_builder::buildClassicExecutableTree( + expCtx()->opCtx, collection(), *_canonicalQuery, *solutions[0], _ws); _children.emplace_back(std::move(newRoot)); _replannedQs = std::move(solutions.back()); solutions.pop_back(); @@ -242,8 +244,7 @@ Status CachedPlanStage::replan(PlanYieldPolicy* yieldPolicy, bool shouldCache, s // Many solutions. Create a MultiPlanStage to pick the best, update the cache, // and so on. The working set will be shared by all candidate plans. - auto cachingMode = shouldCache ? MultiPlanStage::CachingMode::AlwaysCache - : MultiPlanStage::CachingMode::NeverCache; + auto cachingMode = shouldCache ? PlanCachingMode::AlwaysCache : PlanCachingMode::NeverCache; _children.emplace_back( new MultiPlanStage(expCtx(), collection(), _canonicalQuery, cachingMode)); MultiPlanStage* multiPlanStage = static_cast<MultiPlanStage*>(child().get()); @@ -253,8 +254,8 @@ Status CachedPlanStage::replan(PlanYieldPolicy* yieldPolicy, bool shouldCache, s solutions[ix]->cacheData->indexFilterApplied = _plannerParams.indexFiltersApplied; } - auto nextPlanRoot = - StageBuilder::build(opCtx(), collection(), *_canonicalQuery, *solutions[ix], _ws); + auto&& nextPlanRoot = stage_builder::buildClassicExecutableTree( + expCtx()->opCtx, collection(), *_canonicalQuery, *solutions[ix], _ws); multiPlanStage->addPlan(std::move(solutions[ix]), std::move(nextPlanRoot), _ws); } @@ -310,5 +311,4 @@ std::unique_ptr<PlanStageStats> CachedPlanStage::getStats() { const SpecificStats* CachedPlanStage::getSpecificStats() const { return &_specificStats; } - } // namespace mongo diff --git a/src/mongo/db/exec/geo_near.cpp b/src/mongo/db/exec/geo_near.cpp index 10f1ff66888..76290e52fe2 100644 --- a/src/mongo/db/exec/geo_near.cpp +++ b/src/mongo/db/exec/geo_near.cpp @@ -473,62 +473,6 @@ GeoNear2DStage::GeoNear2DStage(const GeoNearParams& nearParams, namespace { - -/** - * Expression which checks whether a legacy 2D index point is contained within our near - * search annulus. See nextInterval() below for more discussion. - * TODO: Make this a standard type of GEO match expression - */ -class TwoDPtInAnnulusExpression : public LeafMatchExpression { -public: - TwoDPtInAnnulusExpression(const R2Annulus& annulus, StringData twoDPath) - : LeafMatchExpression(INTERNAL_2D_POINT_IN_ANNULUS, twoDPath), _annulus(annulus) {} - - void serialize(BSONObjBuilder* out, bool includePath) const final { - out->append("TwoDPtInAnnulusExpression", true); - } - - bool matchesSingleElement(const BSONElement& e, MatchDetails* details = nullptr) const final { - if (!e.isABSONObj()) - return false; - - PointWithCRS point; - if (!GeoParser::parseStoredPoint(e, &point).isOK()) - return false; - - return _annulus.contains(point.oldPoint); - } - - // - // These won't be called. - // - - BSONObj getSerializedRightHandSide() const final { - MONGO_UNREACHABLE; - } - - void debugString(StringBuilder& debug, int level = 0) const final { - MONGO_UNREACHABLE; - } - - bool equivalent(const MatchExpression* other) const final { - MONGO_UNREACHABLE; - return false; - } - - unique_ptr<MatchExpression> shallowClone() const final { - MONGO_UNREACHABLE; - return nullptr; - } - -private: - ExpressionOptimizerFunc getOptimizer() const final { - return [](std::unique_ptr<MatchExpression> expression) { return expression; }; - } - - R2Annulus _annulus; -}; - // Helper class to maintain ownership of a match expression alongside an index scan class FetchStageWithMatch final : public FetchStage { public: diff --git a/src/mongo/db/exec/idhack.cpp b/src/mongo/db/exec/idhack.cpp index f8992a49ff5..196993562d2 100644 --- a/src/mongo/db/exec/idhack.cpp +++ b/src/mongo/db/exec/idhack.cpp @@ -162,16 +162,6 @@ void IDHackStage::doReattachToOperationContext() { _recordCursor->reattachToOperationContext(opCtx()); } -// static -bool IDHackStage::supportsQuery(Collection* collection, const CanonicalQuery& query) { - return !query.getQueryRequest().showRecordId() && query.getQueryRequest().getHint().isEmpty() && - query.getQueryRequest().getMin().isEmpty() && query.getQueryRequest().getMax().isEmpty() && - !query.getQueryRequest().getSkip() && - CanonicalQuery::isSimpleIdQuery(query.getQueryRequest().getFilter()) && - !query.getQueryRequest().isTailable() && - CollatorInterface::collatorsMatch(query.getCollator(), collection->getDefaultCollator()); -} - unique_ptr<PlanStageStats> IDHackStage::getStats() { _commonStats.isEOF = isEOF(); unique_ptr<PlanStageStats> ret = std::make_unique<PlanStageStats>(_commonStats, STAGE_IDHACK); diff --git a/src/mongo/db/exec/idhack.h b/src/mongo/db/exec/idhack.h index edb5b88e449..a06366459e5 100644 --- a/src/mongo/db/exec/idhack.h +++ b/src/mongo/db/exec/idhack.h @@ -68,11 +68,6 @@ public: void doDetachFromOperationContext() final; void doReattachToOperationContext() final; - /** - * ID Hack has a very strict criteria for the queries it supports. - */ - static bool supportsQuery(Collection* collection, const CanonicalQuery& query); - StageType stageType() const final { return STAGE_IDHACK; } diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp index 2c0a28b92bd..2ac307edfac 100644 --- a/src/mongo/db/exec/multi_plan.cpp +++ b/src/mongo/db/exec/multi_plan.cpp @@ -43,6 +43,7 @@ #include "mongo/db/client.h" #include "mongo/db/concurrency/write_conflict_exception.h" #include "mongo/db/exec/scoped_timer.h" +#include "mongo/db/exec/trial_period_utils.h" #include "mongo/db/exec/working_set_common.h" #include "mongo/db/query/collection_query_info.h" #include "mongo/db/query/explain.h" @@ -73,7 +74,7 @@ void markShouldCollectTimingInfoOnSubtree(PlanStage* root) { MultiPlanStage::MultiPlanStage(ExpressionContext* expCtx, const Collection* collection, CanonicalQuery* cq, - CachingMode cachingMode) + PlanCachingMode cachingMode) : RequiresCollectionStage(kStageType, expCtx, collection), _cachingMode(cachingMode), _query(cq), @@ -84,7 +85,7 @@ void MultiPlanStage::addPlan(std::unique_ptr<QuerySolution> solution, std::unique_ptr<PlanStage> root, WorkingSet* ws) { _children.emplace_back(std::move(root)); - _candidates.push_back(CandidatePlan(std::move(solution), _children.back().get(), ws)); + _candidates.push_back({std::move(solution), _children.back().get(), ws}); // Tell the new candidate plan that it must collect timing info. This timing info will // later be stored in the plan cache, and may be used for explain output. @@ -100,12 +101,12 @@ bool MultiPlanStage::isEOF() { // We must have returned all our cached results // and there must be no more results from the best plan. - CandidatePlan& bestPlan = _candidates[_bestPlanIdx]; + auto& bestPlan = _candidates[_bestPlanIdx]; return bestPlan.results.empty() && bestPlan.root->isEOF(); } PlanStage::StageState MultiPlanStage::doWork(WorkingSetID* out) { - CandidatePlan& bestPlan = _candidates[_bestPlanIdx]; + auto& bestPlan = _candidates[_bestPlanIdx]; // Look for an already produced result that provides the data the caller wants. if (!bestPlan.results.empty()) { @@ -151,51 +152,19 @@ void MultiPlanStage::tryYield(PlanYieldPolicy* yieldPolicy) { // 2) some stage requested a yield, or // 3) we need to yield and retry due to a WriteConflictException. // In all cases, the actual yielding happens here. - if (yieldPolicy->shouldYieldOrInterrupt()) { - uassertStatusOK(yieldPolicy->yieldOrInterrupt()); + if (yieldPolicy->shouldYieldOrInterrupt(expCtx()->opCtx)) { + uassertStatusOK(yieldPolicy->yieldOrInterrupt(expCtx()->opCtx)); } } -// static -size_t MultiPlanStage::getTrialPeriodWorks(OperationContext* opCtx, const Collection* collection) { - // Run each plan some number of times. This number is at least as great as - // 'internalQueryPlanEvaluationWorks', but may be larger for big collections. - size_t numWorks = internalQueryPlanEvaluationWorks.load(); - if (nullptr != collection) { - // For large collections, the number of works is set to be this - // fraction of the collection size. - double fraction = internalQueryPlanEvaluationCollFraction; - - numWorks = std::max(static_cast<size_t>(internalQueryPlanEvaluationWorks.load()), - static_cast<size_t>(fraction * collection->numRecords(opCtx))); - } - - return numWorks; -} - -// static -size_t MultiPlanStage::getTrialPeriodNumToReturn(const CanonicalQuery& query) { - // Determine the number of results which we will produce during the plan - // ranking phase before stopping. - size_t numResults = static_cast<size_t>(internalQueryPlanEvaluationMaxResults.load()); - if (query.getQueryRequest().getNToReturn()) { - numResults = - std::min(static_cast<size_t>(*query.getQueryRequest().getNToReturn()), numResults); - } else if (query.getQueryRequest().getLimit()) { - numResults = std::min(static_cast<size_t>(*query.getQueryRequest().getLimit()), numResults); - } - - return numResults; -} - Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { // Adds the amount of time taken by pickBestPlan() to executionTimeMillis. There's lots of // execution work that happens here, so this is needed for the time accounting to // make sense. auto optTimer = getOptTimer(); - size_t numWorks = getTrialPeriodWorks(opCtx(), collection()); - size_t numResults = getTrialPeriodNumToReturn(*_query); + size_t numWorks = trial_period::getTrialPeriodMaxWorks(opCtx(), collection()); + size_t numResults = trial_period::getTrialPeriodNumToReturn(*_query); try { // Work the plans, stopping when a plan hits EOF or returns some fixed number of results. @@ -209,9 +178,9 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { return e.toStatus().withContext("error while multiplanner was selecting best plan"); } - // After picking best plan, ranking will own plan stats from - // candidate solutions (winner and losers). - auto statusWithRanking = PlanRanker::pickBestPlan(_candidates); + // After picking best plan, ranking will own plan stats from candidate solutions (winner and + // losers). + auto statusWithRanking = plan_ranker::pickBestPlan<PlanStageStats>(_candidates); if (!statusWithRanking.isOK()) { return statusWithRanking.getStatus(); } @@ -224,12 +193,7 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { verify(_bestPlanIdx >= 0 && _bestPlanIdx < static_cast<int>(_candidates.size())); - // Copy candidate order and failed candidates. We will need this to sort candidate stats for - // explain after transferring ownership of 'ranking' to plan cache. - std::vector<size_t> candidateOrder = ranking->candidateOrder; - std::vector<size_t> failedCandidates = ranking->failedCandidates; - - CandidatePlan& bestCandidate = _candidates[_bestPlanIdx]; + auto& bestCandidate = _candidates[_bestPlanIdx]; const auto& alreadyProduced = bestCandidate.results; const auto& bestSolution = bestCandidate.solution; @@ -246,7 +210,7 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { _backupPlanIdx = kNoSuchPlan; if (bestSolution->hasBlockingStage && (0 == alreadyProduced.size())) { LOGV2_DEBUG(20592, 5, "Winner has blocking stage, looking for backup plan..."); - for (auto&& ix : candidateOrder) { + for (auto&& ix : ranking->candidateOrder) { if (!_candidates[ix].solution->hasBlockingStage) { LOGV2_DEBUG(20593, 5, "Candidate {ix} is backup child", "ix"_attr = ix); _backupPlanIdx = ix; @@ -255,103 +219,8 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { } } - // Even if the query is of a cacheable shape, the caller might have indicated that we shouldn't - // write to the plan cache. - // - // TODO: We can remove this if we introduce replanning logic to the SubplanStage. - bool canCache = (_cachingMode == CachingMode::AlwaysCache); - if (_cachingMode == CachingMode::SometimesCache) { - // In "sometimes cache" mode, we cache unless we hit one of the special cases below. - canCache = true; - - if (ranking->tieForBest) { - // The winning plan tied with the runner-up and we're using "sometimes cache" mode. We - // will not write a plan cache entry. - canCache = false; - - // These arrays having two or more entries is implied by 'tieForBest'. - invariant(ranking->scores.size() > 1U); - invariant(ranking->candidateOrder.size() > 1U); - - size_t winnerIdx = ranking->candidateOrder[0]; - size_t runnerUpIdx = ranking->candidateOrder[1]; - - LOGV2_DEBUG(20594, - 1, - "Winning plan tied with runner-up. Not caching. query: {query_Short} " - "winner score: {ranking_scores_0} winner summary: " - "{Explain_getPlanSummary_candidates_winnerIdx_root} runner-up score: " - "{ranking_scores_1} runner-up summary: " - "{Explain_getPlanSummary_candidates_runnerUpIdx_root}", - "query_Short"_attr = redact(_query->toStringShort()), - "ranking_scores_0"_attr = ranking->scores[0], - "Explain_getPlanSummary_candidates_winnerIdx_root"_attr = - Explain::getPlanSummary(_candidates[winnerIdx].root), - "ranking_scores_1"_attr = ranking->scores[1], - "Explain_getPlanSummary_candidates_runnerUpIdx_root"_attr = - Explain::getPlanSummary(_candidates[runnerUpIdx].root)); - } - - if (alreadyProduced.empty()) { - // We're using the "sometimes cache" mode, and the winning plan produced no results - // during the plan ranking trial period. We will not write a plan cache entry. - canCache = false; - - size_t winnerIdx = ranking->candidateOrder[0]; - LOGV2_DEBUG(20595, - 1, - "Winning plan had zero results. Not caching. query: {query_Short} winner " - "score: {ranking_scores_0} winner summary: " - "{Explain_getPlanSummary_candidates_winnerIdx_root}", - "query_Short"_attr = redact(_query->toStringShort()), - "ranking_scores_0"_attr = ranking->scores[0], - "Explain_getPlanSummary_candidates_winnerIdx_root"_attr = - Explain::getPlanSummary(_candidates[winnerIdx].root)); - } - } - - // Store the choice we just made in the cache, if the query is of a type that is safe to - // cache. - if (PlanCache::shouldCacheQuery(*_query) && canCache) { - // Create list of candidate solutions for the cache with - // the best solution at the front. - std::vector<QuerySolution*> solutions; - - // Generate solutions and ranking decisions sorted by score. - for (auto&& ix : candidateOrder) { - solutions.push_back(_candidates[ix].solution.get()); - } - // Insert the failed plans in the back. - for (auto&& ix : failedCandidates) { - solutions.push_back(_candidates[ix].solution.get()); - } - - // Check solution cache data. Do not add to cache if - // we have any invalid SolutionCacheData data. - // XXX: One known example is 2D queries - bool validSolutions = true; - for (size_t ix = 0; ix < solutions.size(); ++ix) { - if (nullptr == solutions[ix]->cacheData.get()) { - LOGV2_DEBUG( - 20596, - 5, - "Not caching query because this solution has no cache data: {solutions_ix}", - "solutions_ix"_attr = redact(solutions[ix]->toString())); - validSolutions = false; - break; - } - } - - if (validSolutions) { - CollectionQueryInfo::get(collection()) - .getPlanCache() - ->set(*_query, - solutions, - std::move(ranking), - opCtx()->getServiceContext()->getPreciseClockSource()->now()) - .transitional_ignore(); - } - } + plan_cache_util::updatePlanCache( + expCtx()->opCtx, collection(), _cachingMode, *_query, std::move(ranking), _candidates); return Status::OK(); } @@ -360,7 +229,7 @@ bool MultiPlanStage::workAllPlans(size_t numResults, PlanYieldPolicy* yieldPolic bool doneWorking = false; for (size_t ix = 0; ix < _candidates.size(); ++ix) { - CandidatePlan& candidate = _candidates[ix]; + auto& candidate = _candidates[ix]; if (candidate.failed) { continue; } @@ -391,7 +260,7 @@ bool MultiPlanStage::workAllPlans(size_t numResults, PlanYieldPolicy* yieldPolic if (PlanStage::ADVANCED == state) { // Save result for later. - WorkingSetMember* member = candidate.ws->get(id); + WorkingSetMember* member = candidate.data->get(id); // Ensure that the BSONObj underlying the WorkingSetMember is owned in case we choose to // return the results from the 'candidate' plan. member->makeObjOwnedIfNeeded(); @@ -434,13 +303,20 @@ int MultiPlanStage::bestPlanIdx() const { return _bestPlanIdx; } -QuerySolution* MultiPlanStage::bestSolution() { +const QuerySolution* MultiPlanStage::bestSolution() const { if (_bestPlanIdx == kNoSuchPlan) return nullptr; return _candidates[_bestPlanIdx].solution.get(); } +std::unique_ptr<QuerySolution> MultiPlanStage::bestSolution() { + if (_bestPlanIdx == kNoSuchPlan) + return nullptr; + + return std::move(_candidates[_bestPlanIdx].solution); +} + unique_ptr<PlanStageStats> MultiPlanStage::getStats() { _commonStats.isEOF = isEOF(); unique_ptr<PlanStageStats> ret = diff --git a/src/mongo/db/exec/multi_plan.h b/src/mongo/db/exec/multi_plan.h index 4b9b64bf863..59791dc7d2d 100644 --- a/src/mongo/db/exec/multi_plan.h +++ b/src/mongo/db/exec/multi_plan.h @@ -31,6 +31,7 @@ #include "mongo/db/catalog/collection.h" +#include "mongo/db/exec/plan_cache_util.h" #include "mongo/db/exec/requires_collection_stage.h" #include "mongo/db/exec/working_set.h" #include "mongo/db/jsobj.h" @@ -53,24 +54,6 @@ namespace mongo { class MultiPlanStage final : public RequiresCollectionStage { public: /** - * Callers use this to specify how the MultiPlanStage should interact with the plan cache. - */ - enum class CachingMode { - // Always write a cache entry for the winning plan to the plan cache, overwriting any - // previously existing cache entry for the query shape. - AlwaysCache, - - // Write a cache entry for the query shape *unless* we encounter one of the following edge - // cases: - // - Two or more plans tied for the win. - // - The winning plan returned zero query results during the plan ranking trial period. - SometimesCache, - - // Do not write to the plan cache. - NeverCache, - }; - - /** * Takes no ownership. * * If 'shouldCache' is true, writes a cache entry for the winning plan to the plan cache @@ -79,7 +62,7 @@ public: MultiPlanStage(ExpressionContext* expCtx, const Collection* collection, CanonicalQuery* cq, - CachingMode cachingMode = CachingMode::AlwaysCache); + PlanCachingMode cachingMode = PlanCachingMode::AlwaysCache); bool isEOF() final; @@ -114,19 +97,6 @@ public: */ Status pickBestPlan(PlanYieldPolicy* yieldPolicy); - /** - * Returns the number of times that we are willing to work a plan during a trial period. - * - * Calculated based on a fixed query knob and the size of the collection. - */ - static size_t getTrialPeriodWorks(OperationContext* opCtx, const Collection* collection); - - /** - * Returns the max number of documents which we should allow any plan to return during the - * trial period. As soon as any plan hits this number of documents, the trial period ends. - */ - static size_t getTrialPeriodNumToReturn(const CanonicalQuery& query); - /** Return true if a best plan has been chosen */ bool bestPlanChosen() const; @@ -134,12 +104,20 @@ public: int bestPlanIdx() const; /** - * Returns the QuerySolution for the best plan, or NULL if no best plan + * Returns the QuerySolution for the best plan, or NULL if no best plan. * * The MultiPlanStage retains ownership of the winning QuerySolution and returns an * unowned pointer. */ - QuerySolution* bestSolution(); + const QuerySolution* bestSolution() const; + + /** + * Returns the QuerySolution for the best plan, or NULL if no best plan. + * + * The MultiPlanStage does not retain ownership of the winning QuerySolution and returns + * a unique pointer. + */ + std::unique_ptr<QuerySolution> bestSolution(); /** * Returns true if a backup plan was picked. @@ -185,7 +163,7 @@ private: static const int kNoSuchPlan = -1; // Describes the cases in which we should write an entry for the winning plan to the plan cache. - const CachingMode _cachingMode; + const PlanCachingMode _cachingMode; // The query that we're trying to figure out the best solution to. // not owned here @@ -195,7 +173,7 @@ private: // of all QuerySolutions is retained here, and will *not* be tranferred to the PlanExecutor that // wraps this stage. Ownership of the PlanStages will be in PlanStage::_children which maps // one-to-one with _candidates. - std::vector<CandidatePlan> _candidates; + std::vector<plan_ranker::CandidatePlan> _candidates; // index into _candidates, of the winner of the plan competition // uses -1 / kNoSuchPlan when best plan is not (yet) known diff --git a/src/mongo/db/exec/plan_cache_util.cpp b/src/mongo/db/exec/plan_cache_util.cpp new file mode 100644 index 00000000000..0b40c431921 --- /dev/null +++ b/src/mongo/db/exec/plan_cache_util.cpp @@ -0,0 +1,72 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/plan_cache_util.h" + +#include "mongo/db/query/explain.h" +#include "mongo/logv2/log.h" + +namespace mongo::plan_cache_util { +namespace log_detail { +void logTieForBest(std::string&& query, + double winnerScore, + double runnerUpScore, + std::string winnerPlanSummary, + std::string runnerUpPlanSummary) { + LOGV2_DEBUG(20594, + 1, + "Winning plan tied with runner-up, skip caching", + "query"_attr = redact(query), + "winnerScore"_attr = winnerScore, + "winnerPlanSummary"_attr = winnerPlanSummary, + "runnerUpScore"_attr = runnerUpScore, + "runnerUpPlanSummary"_attr = runnerUpPlanSummary); +} + +void logNotCachingZeroResults(std::string&& query, double score, std::string winnerPlanSummary) { + LOGV2_DEBUG(20595, + 1, + "Winning plan had zero results, skip caching", + "query"_attr = redact(query), + "winnerScore"_attr = score, + "winnerPlanSummary"_attr = winnerPlanSummary); +} + +void logNotCachingNoData(std::string&& solution) { + LOGV2_DEBUG(20596, + 5, + "Not caching query because this solution has no cache data", + "solutions"_attr = redact(solution)); +} +} // namespace log_detail +} // namespace mongo::plan_cache_util diff --git a/src/mongo/db/exec/plan_cache_util.h b/src/mongo/db/exec/plan_cache_util.h new file mode 100644 index 00000000000..3bb25315724 --- /dev/null +++ b/src/mongo/db/exec/plan_cache_util.h @@ -0,0 +1,168 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/plan_stats.h" +#include "mongo/db/query/canonical_query.h" +#include "mongo/db/query/collection_query_info.h" +#include "mongo/db/query/sbe_plan_ranker.h" + +namespace mongo { +/** + * Specifies how the multi-planner should interact with the plan cache. + */ +enum class PlanCachingMode { + // Always write a cache entry for the winning plan to the plan cache, overwriting any + // previously existing cache entry for the query shape. + AlwaysCache, + + // Write a cache entry for the query shape *unless* we encounter one of the following edge + // cases: + // - Two or more plans tied for the win. + // - The winning plan returned zero query results during the plan ranking trial period. + SometimesCache, + + // Do not write to the plan cache. + NeverCache, +}; + +namespace plan_cache_util { +// The logging facility enforces the rule that logging should not be done in a header file. Since +// the template classes and functions below must be defined in the header file and since they do use +// the logging facility, we have to define the helper functions below to perform the actual logging +// operation from template code. +namespace log_detail { +void logTieForBest(std::string&& query, + double winnerScore, + double runnerUpScore, + std::string winnerPlanSummary, + std::string runnerUpPlanSummary); +void logNotCachingZeroResults(std::string&& query, double score, std::string winnerPlanSummary); +void logNotCachingNoData(std::string&& solution); +} // namespace log_detail + +/** + * Caches the best candidate plan, chosen from the given 'candidates' based on the 'ranking' + * decision, if the 'query' is of a type that can be cached. Otherwise, does nothing. + * + * The 'cachingMode' specifies whether the query should be: + * * Always cached. + * * Never cached. + * * Cached, except in certain special cases. + */ +template <typename PlanStageType, typename ResultType, typename Data> +void updatePlanCache( + OperationContext* opCtx, + const Collection* collection, + PlanCachingMode cachingMode, + const CanonicalQuery& query, + std::unique_ptr<plan_ranker::PlanRankingDecision> ranking, + const std::vector<plan_ranker::BaseCandidatePlan<PlanStageType, ResultType, Data>>& + candidates) { + auto winnerIdx = ranking->candidateOrder[0]; + invariant(winnerIdx >= 0 && winnerIdx < candidates.size()); + + // Even if the query is of a cacheable shape, the caller might have indicated that we shouldn't + // write to the plan cache. + // + // TODO: We can remove this if we introduce replanning logic to the SubplanStage. + bool canCache = (cachingMode == PlanCachingMode::AlwaysCache); + if (cachingMode == PlanCachingMode::SometimesCache) { + // In "sometimes cache" mode, we cache unless we hit one of the special cases below. + canCache = true; + + if (ranking->tieForBest) { + // The winning plan tied with the runner-up and we're using "sometimes cache" mode. We + // will not write a plan cache entry. + canCache = false; + + // These arrays having two or more entries is implied by 'tieForBest'. + invariant(ranking->scores.size() > 1U); + invariant(ranking->candidateOrder.size() > 1U); + + size_t winnerIdx = ranking->candidateOrder[0]; + size_t runnerUpIdx = ranking->candidateOrder[1]; + + log_detail::logTieForBest(query.toStringShort(), + ranking->scores[0], + ranking->scores[1], + Explain::getPlanSummary(&*candidates[winnerIdx].root), + Explain::getPlanSummary(&*candidates[runnerUpIdx].root)); + } + + if (candidates[winnerIdx].results.empty()) { + // We're using the "sometimes cache" mode, and the winning plan produced no results + // during the plan ranking trial period. We will not write a plan cache entry. + canCache = false; + log_detail::logNotCachingZeroResults( + query.toStringShort(), + ranking->scores[0], + Explain::getPlanSummary(&*candidates[winnerIdx].root)); + } + } + + // Store the choice we just made in the cache, if the query is of a type that is safe to + // cache. + if (PlanCache::shouldCacheQuery(query) && canCache) { + // Create list of candidate solutions for the cache with the best solution at the front. + std::vector<QuerySolution*> solutions; + + // Generate solutions and ranking decisions sorted by score. + for (auto&& ix : ranking->candidateOrder) { + solutions.push_back(candidates[ix].solution.get()); + } + // Insert the failed plans in the back. + for (auto&& ix : ranking->failedCandidates) { + solutions.push_back(candidates[ix].solution.get()); + } + + // Check solution cache data. Do not add to cache if we have any invalid SolutionCacheData + // data. One known example is 2D queries. + bool validSolutions = true; + for (size_t ix = 0; ix < solutions.size(); ++ix) { + if (nullptr == solutions[ix]->cacheData.get()) { + log_detail::logNotCachingNoData(solutions[ix]->toString()); + validSolutions = false; + break; + } + } + + if (validSolutions) { + uassertStatusOK(CollectionQueryInfo::get(collection) + .getPlanCache() + ->set(query, + solutions, + std::move(ranking), + opCtx->getServiceContext()->getPreciseClockSource()->now())); + } + } +} +} // namespace plan_cache_util +} // namespace mongo diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h index 8e3d2dfadff..3d3137029c8 100644 --- a/src/mongo/db/exec/plan_stats.h +++ b/src/mongo/db/exec/plan_stats.h @@ -112,14 +112,15 @@ struct CommonStats { }; // The universal container for a stage's stats. -struct PlanStageStats { - PlanStageStats(const CommonStats& c, StageType t) : stageType(t), common(c) {} +template <typename C, typename T = void*> +struct BasePlanStageStats { + BasePlanStageStats(const C& c, T t = {}) : stageType(t), common(c) {} /** * Make a deep copy. */ - PlanStageStats* clone() const { - PlanStageStats* stats = new PlanStageStats(common, stageType); + BasePlanStageStats<C, T>* clone() const { + auto stats = new BasePlanStageStats<C, T>(common, stageType); if (specific.get()) { stats->specific.reset(specific->clone()); } @@ -144,23 +145,24 @@ struct PlanStageStats { sizeof(*this); } - // See query/stage_type.h - StageType stageType; + T stageType; // Stats exported by implementing the PlanStage interface. - CommonStats common; + C common; // Per-stage place to stash additional information std::unique_ptr<SpecificStats> specific; // The stats of the node's children. - std::vector<std::unique_ptr<PlanStageStats>> children; + std::vector<std::unique_ptr<BasePlanStageStats<C, T>>> children; private: - PlanStageStats(const PlanStageStats&) = delete; - PlanStageStats& operator=(const PlanStageStats&) = delete; + BasePlanStageStats(const BasePlanStageStats<C, T>&) = delete; + BasePlanStageStats& operator=(const BasePlanStageStats<C, T>&) = delete; }; +using PlanStageStats = BasePlanStageStats<CommonStats, StageType>; + struct AndHashStats : public SpecificStats { AndHashStats() = default; diff --git a/src/mongo/db/exec/projection_executor_builder.cpp b/src/mongo/db/exec/projection_executor_builder.cpp index cd0e3cb49b4..1b712685ea9 100644 --- a/src/mongo/db/exec/projection_executor_builder.cpp +++ b/src/mongo/db/exec/projection_executor_builder.cpp @@ -36,7 +36,7 @@ #include "mongo/db/exec/inclusion_projection_executor.h" #include "mongo/db/pipeline/expression_find_internal.h" #include "mongo/db/query/projection_ast_path_tracking_visitor.h" -#include "mongo/db/query/projection_ast_walker.h" +#include "mongo/db/query/tree_walker.h" #include "mongo/db/query/util/make_data_structure.h" namespace mongo::projection_executor { @@ -254,7 +254,7 @@ auto buildProjectionExecutor(boost::intrusive_ptr<ExpressionContext> expCtx, {std::make_unique<Executor>(expCtx, policies, params[kAllowFastPath]), expCtx}}; ProjectionExecutorVisitor<Executor> executorVisitor{&context}; projection_ast::PathTrackingWalker walker{&context, {&executorVisitor}, {}}; - projection_ast_walker::walk(&walker, root); + tree_walker::walk<true, projection_ast::ASTNode>(root, &walker); if (params[kOptimizeExecutor]) { context.data().executor->optimize(); } diff --git a/src/mongo/db/exec/projection_executor_builder.h b/src/mongo/db/exec/projection_executor_builder.h index b1f1c706bb8..2df36e8eb94 100644 --- a/src/mongo/db/exec/projection_executor_builder.h +++ b/src/mongo/db/exec/projection_executor_builder.h @@ -33,7 +33,6 @@ #include "mongo/db/exec/projection_executor.h" #include "mongo/db/query/projection_ast.h" -#include "mongo/db/query/projection_ast_walker.h" namespace mongo::projection_executor { /** diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript new file mode 100644 index 00000000000..de55c08ad74 --- /dev/null +++ b/src/mongo/db/exec/sbe/SConscript @@ -0,0 +1,78 @@ +# -*- mode: python -*- + +Import("env") + +env.Library( + target='query_sbe_plan_stats', + source=[ + 'stages/plan_stats.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + ] + ) + +env.Library( + target='query_sbe', + source=[ + 'expressions/expression.cpp', + 'stages/branch.cpp', + 'stages/bson_scan.cpp', + 'stages/check_bounds.cpp', + 'stages/co_scan.cpp', + 'stages/exchange.cpp', + 'stages/hash_agg.cpp', + 'stages/hash_join.cpp', + 'stages/limit_skip.cpp', + 'stages/loop_join.cpp', + 'stages/makeobj.cpp', + 'stages/project.cpp', + 'stages/sort.cpp', + 'stages/spool.cpp', + 'stages/stages.cpp', + 'stages/text_match.cpp', + 'stages/traverse.cpp', + 'stages/union.cpp', + 'stages/unwind.cpp', + 'util/debug_print.cpp', + 'values/bson.cpp', + 'values/value.cpp', + 'vm/arith.cpp', + 'vm/vm.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/db/exec/scoped_timer', + '$BUILD_DIR/mongo/db/query/plan_yield_policy', + '$BUILD_DIR/mongo/db/query/query_planner', + '$BUILD_DIR/mongo/db/service_context', + '$BUILD_DIR/mongo/db/storage/index_entry_comparison', + '$BUILD_DIR/mongo/db/storage/key_string', + '$BUILD_DIR/mongo/util/concurrency/thread_pool', + 'query_sbe_plan_stats' + ] + ) + +env.Library( + target='query_sbe_storage', + source=[ + 'stages/ix_scan.cpp', + 'stages/scan.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/db/db_raii', + 'query_sbe' + ] + ) + +env.CppUnitTest( + target='db_sbe_test', + source=[ + 'sbe_test.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/db/concurrency/lock_manager', + '$BUILD_DIR/mongo/unittest/unittest', + 'query_sbe' + ] +) diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp new file mode 100644 index 00000000000..b53d7aae1a2 --- /dev/null +++ b/src/mongo/db/exec/sbe/expressions/expression.cpp @@ -0,0 +1,634 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" + +#include <sstream> + +#include "mongo/db/exec/sbe/stages/spool.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/util/str.h" + +namespace mongo { +namespace sbe { +/** + * This function generates bytecode for testing whether the top of the stack is Nothing. If it is + * not Nothing then code generated by the 'generator' parameter is executed otherwise it is skipped. + * The test is appended to the 'code' parameter. + */ +template <typename F> +std::unique_ptr<vm::CodeFragment> wrapNothingTest(std::unique_ptr<vm::CodeFragment> code, + F&& generator) { + auto inner = std::make_unique<vm::CodeFragment>(); + inner = generator(std::move(inner)); + + invariant(inner->stackSize() == 0); + + // Append the jump that skips around the inner block. + code->appendJumpNothing(inner->instrs().size()); + + code->append(std::move(inner)); + + return code; +} + +std::unique_ptr<EExpression> EConstant::clone() const { + auto [tag, val] = value::copyValue(_tag, _val); + return std::make_unique<EConstant>(tag, val); +} + +std::unique_ptr<vm::CodeFragment> EConstant::compile(CompileCtx& ctx) const { + auto code = std::make_unique<vm::CodeFragment>(); + + code->appendConstVal(_tag, _val); + + return code; +} + +std::vector<DebugPrinter::Block> EConstant::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + std::stringstream ss; + value::printValue(ss, _tag, _val); + + ret.emplace_back(ss.str()); + + return ret; +} + +std::unique_ptr<EExpression> EVariable::clone() const { + return _frameId ? std::make_unique<EVariable>(*_frameId, _var) + : std::make_unique<EVariable>(_var); +} + +std::unique_ptr<vm::CodeFragment> EVariable::compile(CompileCtx& ctx) const { + auto code = std::make_unique<vm::CodeFragment>(); + + if (_frameId) { + int offset = -_var - 1; + code->appendLocalVal(*_frameId, offset); + } else { + auto accessor = ctx.root->getAccessor(ctx, _var); + code->appendAccessVal(accessor); + } + + return code; +} + +std::vector<DebugPrinter::Block> EVariable::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + if (_frameId) { + DebugPrinter::addIdentifier(ret, *_frameId, _var); + } else { + DebugPrinter::addIdentifier(ret, _var); + } + + return ret; +} + +std::unique_ptr<EExpression> EPrimBinary::clone() const { + return std::make_unique<EPrimBinary>(_op, _nodes[0]->clone(), _nodes[1]->clone()); +} + +std::unique_ptr<vm::CodeFragment> EPrimBinary::compile(CompileCtx& ctx) const { + auto code = std::make_unique<vm::CodeFragment>(); + + auto lhs = _nodes[0]->compile(ctx); + auto rhs = _nodes[1]->compile(ctx); + + switch (_op) { + case EPrimBinary::add: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendAdd(); + break; + case EPrimBinary::sub: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendSub(); + break; + case EPrimBinary::mul: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendMul(); + break; + case EPrimBinary::div: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendDiv(); + break; + case EPrimBinary::less: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendLess(); + break; + case EPrimBinary::lessEq: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendLessEq(); + break; + case EPrimBinary::greater: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendGreater(); + break; + case EPrimBinary::greaterEq: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendGreaterEq(); + break; + case EPrimBinary::eq: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendEq(); + break; + case EPrimBinary::neq: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendNeq(); + break; + case EPrimBinary::cmp3w: + code->append(std::move(lhs)); + code->append(std::move(rhs)); + code->appendCmp3w(); + break; + case EPrimBinary::logicAnd: { + auto codeFalseBranch = std::make_unique<vm::CodeFragment>(); + codeFalseBranch->appendConstVal(value::TypeTags::Boolean, false); + // Jump to the merge point that will be right after the thenBranch (rhs). + codeFalseBranch->appendJump(rhs->instrs().size()); + + code->append(std::move(lhs)); + code = wrapNothingTest(std::move(code), [&](std::unique_ptr<vm::CodeFragment> code) { + code->appendJumpTrue(codeFalseBranch->instrs().size()); + code->append(std::move(codeFalseBranch), std::move(rhs)); + + return code; + }); + break; + } + case EPrimBinary::logicOr: { + auto codeTrueBranch = std::make_unique<vm::CodeFragment>(); + codeTrueBranch->appendConstVal(value::TypeTags::Boolean, true); + + // Jump to the merge point that will be right after the thenBranch (true branch). + rhs->appendJump(codeTrueBranch->instrs().size()); + + code->append(std::move(lhs)); + code = wrapNothingTest(std::move(code), [&](std::unique_ptr<vm::CodeFragment> code) { + code->appendJumpTrue(rhs->instrs().size()); + code->append(std::move(rhs), std::move(codeTrueBranch)); + + return code; + }); + break; + } + default: + MONGO_UNREACHABLE; + break; + } + return code; +} + +std::vector<DebugPrinter::Block> EPrimBinary::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + DebugPrinter::addBlocks(ret, _nodes[0]->debugPrint()); + + switch (_op) { + case EPrimBinary::add: + ret.emplace_back("+"); + break; + case EPrimBinary::sub: + ret.emplace_back("-"); + break; + case EPrimBinary::mul: + ret.emplace_back("*"); + break; + case EPrimBinary::div: + ret.emplace_back("/"); + break; + case EPrimBinary::less: + ret.emplace_back("<"); + break; + case EPrimBinary::lessEq: + ret.emplace_back("<="); + break; + case EPrimBinary::greater: + ret.emplace_back(">"); + break; + case EPrimBinary::greaterEq: + ret.emplace_back(">="); + break; + case EPrimBinary::eq: + ret.emplace_back("=="); + break; + case EPrimBinary::neq: + ret.emplace_back("!="); + break; + case EPrimBinary::cmp3w: + ret.emplace_back("<=>"); + break; + case EPrimBinary::logicAnd: + ret.emplace_back("&&"); + break; + case EPrimBinary::logicOr: + ret.emplace_back("||"); + break; + default: + MONGO_UNREACHABLE; + break; + } + DebugPrinter::addBlocks(ret, _nodes[1]->debugPrint()); + + return ret; +} + +std::unique_ptr<EExpression> EPrimUnary::clone() const { + return std::make_unique<EPrimUnary>(_op, _nodes[0]->clone()); +} + +std::unique_ptr<vm::CodeFragment> EPrimUnary::compile(CompileCtx& ctx) const { + auto code = std::make_unique<vm::CodeFragment>(); + + auto operand = _nodes[0]->compile(ctx); + + switch (_op) { + case negate: + code->append(std::move(operand)); + code->appendNegate(); + break; + case EPrimUnary::logicNot: + code->append(std::move(operand)); + code->appendNot(); + break; + default: + MONGO_UNREACHABLE; + break; + } + return code; +} + +std::vector<DebugPrinter::Block> EPrimUnary::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + switch (_op) { + case EPrimUnary::negate: + ret.emplace_back("-"); + break; + case EPrimUnary::logicNot: + ret.emplace_back("!"); + break; + default: + MONGO_UNREACHABLE; + break; + } + + DebugPrinter::addBlocks(ret, _nodes[0]->debugPrint()); + + return ret; +} + +std::unique_ptr<EExpression> EFunction::clone() const { + std::vector<std::unique_ptr<EExpression>> args; + args.reserve(_nodes.size()); + for (auto& a : _nodes) { + args.emplace_back(a->clone()); + } + return std::make_unique<EFunction>(_name, std::move(args)); +} + +namespace { +/** + * The arity test function. It returns true if the number of arguments is correct. + */ +using ArityFn = bool (*)(size_t); + +/** + * The builtin function description. + */ +struct BuiltinFn { + ArityFn arityTest; + vm::Builtin builtin; + bool aggregate; +}; + +/** + * The map of recognized builtin functions. + */ +static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = { + {"split", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::split, false}}, + {"regexMatch", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::regexMatch, false}}, + {"dropFields", BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::dropFields, false}}, + {"newObj", BuiltinFn{[](size_t n) { return n % 2 == 0; }, vm::Builtin::newObj, false}}, + {"ksToString", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::ksToString, false}}, + {"ks", BuiltinFn{[](size_t n) { return n > 2; }, vm::Builtin::newKs, false}}, + {"abs", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::abs, false}}, + {"addToArray", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::addToArray, true}}, + {"addToSet", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::addToSet, true}}, +}; + +/** + * The code generation function. + */ +using CodeFn = void (vm::CodeFragment::*)(); + +/** + * The function description. + */ +struct InstrFn { + ArityFn arityTest; + CodeFn generate; + bool aggregate; +}; + +/** + * The map of functions that resolve directly to instructions. + */ +static stdx::unordered_map<std::string, InstrFn> kInstrFunctions = { + {"getField", + InstrFn{[](size_t n) { return n == 2; }, &vm::CodeFragment::appendGetField, false}}, + {"fillEmpty", + InstrFn{[](size_t n) { return n == 2; }, &vm::CodeFragment::appendFillEmpty, false}}, + {"exists", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendExists, false}}, + {"isNull", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsNull, false}}, + {"isObject", + InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsObject, false}}, + {"isArray", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsArray, false}}, + {"isString", + InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsString, false}}, + {"isNumber", + InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendIsNumber, false}}, + {"sum", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendSum, true}}, + {"min", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendMin, true}}, + {"max", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendMax, true}}, + {"first", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendFirst, true}}, + {"last", InstrFn{[](size_t n) { return n == 1; }, &vm::CodeFragment::appendLast, true}}, +}; +} // namespace + +std::unique_ptr<vm::CodeFragment> EFunction::compile(CompileCtx& ctx) const { + if (auto it = kBuiltinFunctions.find(_name); it != kBuiltinFunctions.end()) { + auto arity = _nodes.size(); + if (!it->second.arityTest(arity)) { + uasserted(4822843, + str::stream() << "function call: " << _name << " has wrong arity: " << arity); + } + auto code = std::make_unique<vm::CodeFragment>(); + + for (size_t idx = arity; idx-- > 0;) { + code->append(_nodes[idx]->compile(ctx)); + } + + if (it->second.aggregate) { + uassert(4822844, + str::stream() << "aggregate function call: " << _name + << " occurs in the non-aggregate context.", + ctx.aggExpression); + + code->appendMoveVal(ctx.accumulator); + ++arity; + } + + code->appendFunction(it->second.builtin, arity); + + return code; + } + + if (auto it = kInstrFunctions.find(_name); it != kInstrFunctions.end()) { + if (!it->second.arityTest(_nodes.size())) { + uasserted(4822845, + str::stream() + << "function call: " << _name << " has wrong arity: " << _nodes.size()); + } + auto code = std::make_unique<vm::CodeFragment>(); + + if (it->second.aggregate) { + uassert(4822846, + str::stream() << "aggregate function call: " << _name + << " occurs in the non-aggregate context.", + ctx.aggExpression); + + code->appendAccessVal(ctx.accumulator); + } + + // The order of evaluation is flipped for instruction functions. We may want to change the + // evaluation code for those functions so we have the same behavior for all functions. + for (size_t idx = 0; idx < _nodes.size(); ++idx) { + code->append(_nodes[idx]->compile(ctx)); + } + (*code.*(it->second.generate))(); + + return code; + } + + uasserted(4822847, str::stream() << "unknown function call: " << _name); +} + +std::vector<DebugPrinter::Block> EFunction::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, _name); + + ret.emplace_back("(`"); + for (size_t idx = 0; idx < _nodes.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addBlocks(ret, _nodes[idx]->debugPrint()); + } + ret.emplace_back("`)"); + + return ret; +} + +std::unique_ptr<EExpression> EIf::clone() const { + return std::make_unique<EIf>(_nodes[0]->clone(), _nodes[1]->clone(), _nodes[2]->clone()); +} + +std::unique_ptr<vm::CodeFragment> EIf::compile(CompileCtx& ctx) const { + auto code = std::make_unique<vm::CodeFragment>(); + + auto thenBranch = _nodes[1]->compile(ctx); + + auto elseBranch = _nodes[2]->compile(ctx); + + // The then and else branches must be balanced. + invariant(thenBranch->stackSize() == elseBranch->stackSize()); + + // Jump to the merge point that will be right after the thenBranch. + elseBranch->appendJump(thenBranch->instrs().size()); + + // Compile the condition. + code->append(_nodes[0]->compile(ctx)); + code = wrapNothingTest(std::move(code), [&](std::unique_ptr<vm::CodeFragment> code) { + // Jump around the elseBranch. + code->appendJumpTrue(elseBranch->instrs().size()); + // Append else and then branches. + code->append(std::move(elseBranch), std::move(thenBranch)); + + return code; + }); + return code; +} + +std::vector<DebugPrinter::Block> EIf::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "if"); + + ret.emplace_back("(`"); + + // Print the condition. + DebugPrinter::addBlocks(ret, _nodes[0]->debugPrint()); + ret.emplace_back(DebugPrinter::Block("`,")); + // Print thenBranch. + DebugPrinter::addBlocks(ret, _nodes[1]->debugPrint()); + ret.emplace_back(DebugPrinter::Block("`,")); + // Print elseBranch. + DebugPrinter::addBlocks(ret, _nodes[2]->debugPrint()); + + ret.emplace_back("`)"); + + return ret; +} + +std::unique_ptr<EExpression> ELocalBind::clone() const { + std::vector<std::unique_ptr<EExpression>> binds; + binds.reserve(_nodes.size() - 1); + for (size_t idx = 0; idx < _nodes.size() - 1; ++idx) { + binds.emplace_back(_nodes[idx]->clone()); + } + return std::make_unique<ELocalBind>(_frameId, std::move(binds), _nodes.back()->clone()); +} + +std::unique_ptr<vm::CodeFragment> ELocalBind::compile(CompileCtx& ctx) const { + auto code = std::make_unique<vm::CodeFragment>(); + + // Generate bytecode for local variables and the 'in' expression. The 'in' expression is in the + // last position of _nodes. + for (size_t idx = 0; idx < _nodes.size(); ++idx) { + auto c = _nodes[idx]->compile(ctx); + code->append(std::move(c)); + } + + // After the execution we have to cleanup the stack; i.e. local variables go out of scope. + // However, note that the top of the stack holds the overall result (i.e. the 'in' expression) + // and it cannot be destroyed. So we 'bubble' it down with a series of swap/pop instructions. + for (size_t idx = 0; idx < _nodes.size() - 1; ++idx) { + code->appendSwap(); + code->appendPop(); + } + + // Local variables are no longer accessible after this point so remove any fixup information. + code->removeFixup(_frameId); + return code; +} + +std::vector<DebugPrinter::Block> ELocalBind::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + DebugPrinter::addKeyword(ret, "let"); + + ret.emplace_back("[`"); + for (size_t idx = 0; idx < _nodes.size() - 1; ++idx) { + if (idx != 0) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _frameId, idx); + ret.emplace_back("="); + DebugPrinter::addBlocks(ret, _nodes[idx]->debugPrint()); + } + ret.emplace_back("`]"); + + DebugPrinter::addBlocks(ret, _nodes.back()->debugPrint()); + + return ret; +} + +std::unique_ptr<EExpression> EFail::clone() const { + return std::make_unique<EFail>(_code, _message); +} + +std::unique_ptr<vm::CodeFragment> EFail::compile(CompileCtx& ctx) const { + auto code = std::make_unique<vm::CodeFragment>(); + + code->appendConstVal(value::TypeTags::NumberInt64, + value::bitcastFrom(static_cast<int64_t>(_code))); + + code->appendConstVal(value::TypeTags::StringBig, value::bitcastFrom(_message.c_str())); + + code->appendFail(); + + return code; +} + +std::vector<DebugPrinter::Block> EFail::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "fail"); + + ret.emplace_back("("); + + ret.emplace_back(DebugPrinter::Block(std::to_string(_code))); + ret.emplace_back(DebugPrinter::Block(",`")); + ret.emplace_back(DebugPrinter::Block(_message)); + + ret.emplace_back("`)"); + + return ret; +} + +value::SlotAccessor* CompileCtx::getAccessor(value::SlotId slot) { + for (auto it = correlated.rbegin(); it != correlated.rend(); ++it) { + if (it->first == slot) { + return it->second; + } + } + + uasserted(4822848, str::stream() << "undefined slot accessor:" << slot); +} + +std::shared_ptr<SpoolBuffer> CompileCtx::getSpoolBuffer(SpoolId spool) { + if (spoolBuffers.find(spool) == spoolBuffers.end()) { + spoolBuffers.emplace(spool, std::make_shared<SpoolBuffer>()); + } + return spoolBuffers[spool]; +} + +void CompileCtx::pushCorrelated(value::SlotId slot, value::SlotAccessor* accessor) { + correlated.emplace_back(slot, accessor); +} + +void CompileCtx::popCorrelated() { + correlated.pop_back(); +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/expressions/expression.h b/src/mongo/db/exec/sbe/expressions/expression.h new file mode 100644 index 00000000000..774b739a67d --- /dev/null +++ b/src/mongo/db/exec/sbe/expressions/expression.h @@ -0,0 +1,355 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <memory> +#include <string> +#include <vector> + +#include "mongo/db/exec/sbe/util/debug_print.h" +#include "mongo/db/exec/sbe/values/value.h" +#include "mongo/db/exec/sbe/vm/vm.h" +#include "mongo/stdx/unordered_map.h" + +namespace mongo { +namespace sbe { +using SpoolBuffer = std::vector<value::MaterializedRow>; + +class PlanStage; +struct CompileCtx { + value::SlotAccessor* getAccessor(value::SlotId slot); + std::shared_ptr<SpoolBuffer> getSpoolBuffer(SpoolId spool); + + void pushCorrelated(value::SlotId slot, value::SlotAccessor* accessor); + void popCorrelated(); + + PlanStage* root{nullptr}; + value::SlotAccessor* accumulator{nullptr}; + std::vector<std::pair<value::SlotId, value::SlotAccessor*>> correlated; + stdx::unordered_map<SpoolId, std::shared_ptr<SpoolBuffer>> spoolBuffers; + bool aggExpression{false}; +}; + +/** + * This is an abstract base class of all expression types in SBE. The expression types derived form + * this base must implement two fundamental operations: + * - compile method that generates bytecode that is executed by the VM during runtime + * - clone method that creates a complete copy of the expression + * + * The debugPrint method generates textual representation of the expression for internal debugging + * purposes. + */ +class EExpression { +public: + virtual ~EExpression() = default; + + /** + * The idiomatic C++ pattern of object cloning. Expressions must be fully copyable as every + * thread in parallel execution needs its own private copy. + */ + virtual std::unique_ptr<EExpression> clone() const = 0; + + /** + * Returns bytecode directly executable by VM. + */ + virtual std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const = 0; + + virtual std::vector<DebugPrinter::Block> debugPrint() const = 0; + +protected: + std::vector<std::unique_ptr<EExpression>> _nodes; + + /** + * Expressions can never be constructed with nullptr children. + */ + void validateNodes() { + for (auto& node : _nodes) { + invariant(node); + } + } +}; + +template <typename T, typename... Args> +inline std::unique_ptr<EExpression> makeE(Args&&... args) { + return std::make_unique<T>(std::forward<Args>(args)...); +} + +template <typename... Ts> +inline std::vector<std::unique_ptr<EExpression>> makeEs(Ts&&... pack) { + std::vector<std::unique_ptr<EExpression>> exprs; + + (exprs.emplace_back(std::forward<Ts>(pack)), ...); + + return exprs; +} + +namespace detail { +// base case +inline void makeEM_unwind(value::SlotMap<std::unique_ptr<EExpression>>& result, + value::SlotId slot, + std::unique_ptr<EExpression> expr) { + result.emplace(slot, std::move(expr)); +} + +// recursive case +template <typename... Ts> +inline void makeEM_unwind(value::SlotMap<std::unique_ptr<EExpression>>& result, + value::SlotId slot, + std::unique_ptr<EExpression> expr, + Ts&&... rest) { + result.emplace(slot, std::move(expr)); + makeEM_unwind(result, std::forward<Ts>(rest)...); +} +} // namespace detail + +template <typename... Ts> +auto makeEM(Ts&&... pack) { + value::SlotMap<std::unique_ptr<EExpression>> result; + if constexpr (sizeof...(pack) > 0) { + result.reserve(sizeof...(Ts) / 2); + detail::makeEM_unwind(result, std::forward<Ts>(pack)...); + } + return result; +} + +template <typename... Args> +auto makeSV(Args&&... args) { + value::SlotVector v; + v.reserve(sizeof...(Args)); + (v.push_back(std::forward<Args>(args)), ...); + return v; +} + +/** + * This is a constant expression. It assumes the ownership of the input constant. + */ +class EConstant final : public EExpression { +public: + EConstant(value::TypeTags tag, value::Value val) : _tag(tag), _val(val) {} + EConstant(std::string_view str) { + // Views are non-owning so we have to make a copy. + auto [tag, val] = value::makeNewString(str); + + _tag = tag; + _val = val; + } + + ~EConstant() override { + value::releaseValue(_tag, _val); + } + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; + +private: + value::TypeTags _tag; + value::Value _val; +}; + +/** + * This is an expression representing a variable. The variable can point to a slot as defined by a + * SBE plan stages or to a slot defined by a local bind (a.k.a. let) expression. The local binds are + * identified by the frame id. + */ +class EVariable final : public EExpression { +public: + EVariable(value::SlotId var) : _var(var), _frameId(boost::none) {} + EVariable(FrameId frameId, value::SlotId var) : _var(var), _frameId(frameId) {} + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; + +private: + value::SlotId _var; + boost::optional<FrameId> _frameId; +}; + +/** + * This is a binary primitive (builtin) operation. + */ +class EPrimBinary final : public EExpression { +public: + enum Op { + add, + sub, + + mul, + div, + + lessEq, + less, + greater, + greaterEq, + + eq, + neq, + + cmp3w, + + // Logical operations are short - circuiting. + logicAnd, + logicOr, + }; + + EPrimBinary(Op op, std::unique_ptr<EExpression> lhs, std::unique_ptr<EExpression> rhs) + : _op(op) { + _nodes.emplace_back(std::move(lhs)); + _nodes.emplace_back(std::move(rhs)); + validateNodes(); + } + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; + +private: + Op _op; +}; + +/** + * This is a unary primitive (builtin) operation. + */ +class EPrimUnary final : public EExpression { +public: + enum Op { + logicNot, + negate, + }; + + EPrimUnary(Op op, std::unique_ptr<EExpression> operand) : _op(op) { + _nodes.emplace_back(std::move(operand)); + validateNodes(); + } + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; + +private: + Op _op; +}; + +/** + * This is a function call expression. Functions can have arbitrary arity and arguments are + * evaluated right to left. They are identified simply by a name and we have a dictionary of all + * supported (builtin) functions. + */ +class EFunction final : public EExpression { +public: + EFunction(std::string_view name, std::vector<std::unique_ptr<EExpression>> args) : _name(name) { + _nodes = std::move(args); + validateNodes(); + } + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; + +private: + std::string _name; +}; + +/** + * This is a conditional (a.k.a. ite) expression. + */ +class EIf final : public EExpression { +public: + EIf(std::unique_ptr<EExpression> cond, + std::unique_ptr<EExpression> thenBranch, + std::unique_ptr<EExpression> elseBranch) { + _nodes.emplace_back(std::move(cond)); + _nodes.emplace_back(std::move(thenBranch)); + _nodes.emplace_back(std::move(elseBranch)); + validateNodes(); + } + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; +}; + +/** + * This is a let expression that can be used to define local variables. + */ +class ELocalBind final : public EExpression { +public: + ELocalBind(FrameId frameId, + std::vector<std::unique_ptr<EExpression>> binds, + std::unique_ptr<EExpression> in) + : _frameId(frameId) { + _nodes = std::move(binds); + _nodes.emplace_back(std::move(in)); + validateNodes(); + } + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; + +private: + FrameId _frameId; +}; + +/** + * Evaluating this expression will throw an exception with given error code and message. + */ +class EFail final : public EExpression { +public: + EFail(ErrorCodes::Error code, std::string message) + : _code(code), _message(std::move(message)) {} + + std::unique_ptr<EExpression> clone() const override; + + std::unique_ptr<vm::CodeFragment> compile(CompileCtx& ctx) const override; + + std::vector<DebugPrinter::Block> debugPrint() const override; + +private: + ErrorCodes::Error _code; + std::string _message; +}; +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/sbe_test.cpp b/src/mongo/db/exec/sbe/sbe_test.cpp new file mode 100644 index 00000000000..6b06e6229c6 --- /dev/null +++ b/src/mongo/db/exec/sbe/sbe_test.cpp @@ -0,0 +1,162 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/db/exec/sbe/values/value.h" +#include "mongo/db/exec/sbe/vm/vm.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::sbe { + +TEST(SBEValues, Basic) { + using namespace std::literals; + { + const auto [tag, val] = value::makeNewString("small"sv); + ASSERT_EQUALS(tag, value::TypeTags::StringSmall); + + value::releaseValue(tag, val); + } + + { + const auto [tag, val] = value::makeNewString("not so small string"sv); + ASSERT_EQUALS(tag, value::TypeTags::StringBig); + + value::releaseValue(tag, val); + } + { + const auto [tag, val] = value::makeNewObject(); + auto obj = value::getObjectView(val); + + const auto [fieldTag, fieldVal] = value::makeNewString("not so small string"sv); + obj->push_back("field"sv, fieldTag, fieldVal); + + ASSERT_EQUALS(obj->size(), 1); + const auto [checkTag, checkVal] = obj->getField("field"sv); + + ASSERT_EQUALS(fieldTag, checkTag); + ASSERT_EQUALS(fieldVal, checkVal); + + value::releaseValue(tag, val); + } + { + const auto [tag, val] = value::makeNewArray(); + auto obj = value::getArrayView(val); + + const auto [fieldTag, fieldVal] = value::makeNewString("not so small string"sv); + obj->push_back(fieldTag, fieldVal); + + ASSERT_EQUALS(obj->size(), 1); + const auto [checkTag, checkVal] = obj->getAt(0); + + ASSERT_EQUALS(fieldTag, checkTag); + ASSERT_EQUALS(fieldVal, checkVal); + + value::releaseValue(tag, val); + } +} + +TEST(SBEValues, Hash) { + auto tagInt32 = value::TypeTags::NumberInt32; + auto valInt32 = value::bitcastFrom<int32_t>(-5); + + auto tagInt64 = value::TypeTags::NumberInt64; + auto valInt64 = value::bitcastFrom<int64_t>(-5); + + auto tagDouble = value::TypeTags::NumberDouble; + auto valDouble = value::bitcastFrom<double>(-5.0); + + auto [tagDecimal, valDecimal] = value::makeCopyDecimal(mongo::Decimal128(-5.0)); + + ASSERT_EQUALS(value::hashValue(tagInt32, valInt32), value::hashValue(tagInt64, valInt64)); + ASSERT_EQUALS(value::hashValue(tagInt32, valInt32), value::hashValue(tagDouble, valDouble)); + ASSERT_EQUALS(value::hashValue(tagInt32, valInt32), value::hashValue(tagDecimal, valDecimal)); + + value::releaseValue(tagDecimal, valDecimal); +} + +TEST(SBEVM, Add) { + { + auto tagInt32 = value::TypeTags::NumberInt32; + auto valInt32 = value::bitcastFrom<int32_t>(-7); + + auto tagInt64 = value::TypeTags::NumberInt64; + auto valInt64 = value::bitcastFrom<int64_t>(-5); + + vm::CodeFragment code; + code.appendConstVal(tagInt32, valInt32); + code.appendConstVal(tagInt64, valInt64); + code.appendAdd(); + + vm::ByteCode interpreter; + auto [owned, tag, val] = interpreter.run(&code); + + ASSERT_EQUALS(tag, value::TypeTags::NumberInt64); + ASSERT_EQUALS(val, -12); + } + { + auto tagInt32 = value::TypeTags::NumberInt32; + auto valInt32 = value::bitcastFrom<int32_t>(-7); + + auto tagDouble = value::TypeTags::NumberDouble; + auto valDouble = value::bitcastFrom<double>(-5.0); + + vm::CodeFragment code; + code.appendConstVal(tagInt32, valInt32); + code.appendConstVal(tagDouble, valDouble); + code.appendAdd(); + + vm::ByteCode interpreter; + auto [owned, tag, val] = interpreter.run(&code); + + ASSERT_EQUALS(tag, value::TypeTags::NumberDouble); + ASSERT_EQUALS(value::bitcastTo<double>(val), -12.0); + } + { + auto [tagDecimal, valDecimal] = value::makeCopyDecimal(mongo::Decimal128(-7.25)); + + auto tagDouble = value::TypeTags::NumberDouble; + auto valDouble = value::bitcastFrom<double>(-5.25); + + vm::CodeFragment code; + code.appendConstVal(tagDecimal, valDecimal); + code.appendConstVal(tagDouble, valDouble); + code.appendAdd(); + + vm::ByteCode interpreter; + auto [owned, tag, val] = interpreter.run(&code); + + ASSERT_EQUALS(tag, value::TypeTags::NumberDecimal); + ASSERT_EQUALS(value::bitcastTo<mongo::Decimal128>(val).toDouble(), -12.5); + ASSERT_TRUE(owned); + + value::releaseValue(tag, val); + value::releaseValue(tagDecimal, valDecimal); + } +} + +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/branch.cpp b/src/mongo/db/exec/sbe/stages/branch.cpp new file mode 100644 index 00000000000..27c3ab2845e --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/branch.cpp @@ -0,0 +1,227 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/branch.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" + +namespace mongo { +namespace sbe { +BranchStage::BranchStage(std::unique_ptr<PlanStage> inputThen, + std::unique_ptr<PlanStage> inputElse, + std::unique_ptr<EExpression> filter, + value::SlotVector inputThenVals, + value::SlotVector inputElseVals, + value::SlotVector outputVals) + : PlanStage("branch"_sd), + _filter(std::move(filter)), + _inputThenVals(std::move(inputThenVals)), + _inputElseVals(std::move(inputElseVals)), + _outputVals(std::move(outputVals)) { + invariant(_inputThenVals.size() == _outputVals.size()); + invariant(_inputElseVals.size() == _outputVals.size()); + _children.emplace_back(std::move(inputThen)); + _children.emplace_back(std::move(inputElse)); +} + +std::unique_ptr<PlanStage> BranchStage::clone() const { + return std::make_unique<BranchStage>(_children[0]->clone(), + _children[1]->clone(), + _filter->clone(), + _inputThenVals, + _inputElseVals, + _outputVals); +} + +void BranchStage::prepare(CompileCtx& ctx) { + value::SlotSet dupCheck; + + _children[0]->prepare(ctx); + _children[1]->prepare(ctx); + + for (auto slot : _inputThenVals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822829, str::stream() << "duplicate field: " << slot, inserted); + + _inputThenAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + } + + for (auto slot : _inputElseVals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822830, str::stream() << "duplicate field: " << slot, inserted); + + _inputElseAccessors.emplace_back(_children[1]->getAccessor(ctx, slot)); + } + + for (auto slot : _outputVals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822831, str::stream() << "duplicate field: " << slot, inserted); + + _outValueAccessors.emplace_back(value::ViewOfValueAccessor{}); + } + + // compile filter + ctx.root = this; + _filterCode = _filter->compile(ctx); +} + +value::SlotAccessor* BranchStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + for (size_t idx = 0; idx < _outputVals.size(); idx++) { + if (_outputVals[idx] == slot) { + return &_outValueAccessors[idx]; + } + } + + return ctx.getAccessor(slot); +} + +void BranchStage::open(bool reOpen) { + _commonStats.opens++; + _specificStats.numTested++; + + // run the filter expressions here + auto [owned, tag, val] = _bytecode.run(_filterCode.get()); + if (owned) { + value::releaseValue(tag, val); + } + if (tag == value::TypeTags::Boolean) { + if (val) { + _activeBranch = 0; + _children[0]->open(reOpen && _thenOpened); + _thenOpened = true; + } else { + _activeBranch = 1; + _children[1]->open(reOpen && _elseOpened); + _elseOpened = true; + } + } else { + _activeBranch = boost::none; + } +} + +PlanState BranchStage::getNext() { + if (!_activeBranch) { + return trackPlanState(PlanState::IS_EOF); + } + + if (*_activeBranch == 0) { + auto state = _children[0]->getNext(); + if (state == PlanState::ADVANCED) { + for (size_t idx = 0; idx < _outValueAccessors.size(); ++idx) { + auto [tag, val] = _inputThenAccessors[idx]->getViewOfValue(); + _outValueAccessors[idx].reset(tag, val); + } + } + return trackPlanState(state); + } else { + auto state = _children[1]->getNext(); + if (state == PlanState::ADVANCED) { + for (size_t idx = 0; idx < _outValueAccessors.size(); ++idx) { + auto [tag, val] = _inputElseAccessors[idx]->getViewOfValue(); + _outValueAccessors[idx].reset(tag, val); + } + } + return trackPlanState(state); + } +} + +void BranchStage::close() { + _commonStats.closes++; + + if (_thenOpened) { + _children[0]->close(); + _thenOpened = false; + } + if (_elseOpened) { + _children[1]->close(); + _elseOpened = false; + } +} + +std::unique_ptr<PlanStageStats> BranchStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->specific = std::make_unique<FilterStats>(_specificStats); + ret->children.emplace_back(_children[0]->getStats()); + ret->children.emplace_back(_children[1]->getStats()); + return ret; +} + +const SpecificStats* BranchStage::getSpecificStats() const { + return &_specificStats; +} + +std::vector<DebugPrinter::Block> BranchStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "branch"); + + ret.emplace_back("{`"); + DebugPrinter::addBlocks(ret, _filter->debugPrint()); + ret.emplace_back("`}"); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _outputVals.size(); idx++) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + DebugPrinter::addIdentifier(ret, _outputVals[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + DebugPrinter::addNewLine(ret); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _inputThenVals.size(); idx++) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + DebugPrinter::addIdentifier(ret, _inputThenVals[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + DebugPrinter::addNewLine(ret); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _inputElseVals.size(); idx++) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + DebugPrinter::addIdentifier(ret, _inputElseVals[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + DebugPrinter::addBlocks(ret, _children[1]->debugPrint()); + return ret; +} + +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/branch.h b/src/mongo/db/exec/sbe/stages/branch.h new file mode 100644 index 00000000000..9d9005e6d9f --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/branch.h @@ -0,0 +1,80 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/vm/vm.h" + +namespace mongo::sbe { +/** + * This stage delivers results from either 'then' or 'else' branch depending on the value of the + * 'filer' expression as evaluated during the open() call. + */ +class BranchStage final : public PlanStage { +public: + BranchStage(std::unique_ptr<PlanStage> inputThen, + std::unique_ptr<PlanStage> inputElse, + std::unique_ptr<EExpression> filter, + value::SlotVector inputThenVals, + value::SlotVector inputElseVals, + value::SlotVector outputVals); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + const std::unique_ptr<EExpression> _filter; + const value::SlotVector _inputThenVals; + const value::SlotVector _inputElseVals; + const value::SlotVector _outputVals; + std::unique_ptr<vm::CodeFragment> _filterCode; + + std::vector<value::SlotAccessor*> _inputThenAccessors; + std::vector<value::SlotAccessor*> _inputElseAccessors; + std::vector<value::ViewOfValueAccessor> _outValueAccessors; + + boost::optional<int> _activeBranch; + bool _thenOpened{false}; + bool _elseOpened{false}; + + vm::ByteCode _bytecode; + FilterStats _specificStats; +}; +} // namespace mongo::sbe
\ No newline at end of file diff --git a/src/mongo/db/exec/sbe/stages/bson_scan.cpp b/src/mongo/db/exec/sbe/stages/bson_scan.cpp new file mode 100644 index 00000000000..4b5a40b4814 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/bson_scan.cpp @@ -0,0 +1,168 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/bson_scan.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/util/str.h" + +namespace mongo { +namespace sbe { +BSONScanStage::BSONScanStage(const char* bsonBegin, + const char* bsonEnd, + boost::optional<value::SlotId> recordSlot, + std::vector<std::string> fields, + value::SlotVector vars) + : PlanStage("bsonscan"_sd), + _bsonBegin(bsonBegin), + _bsonEnd(bsonEnd), + _recordSlot(recordSlot), + _fields(std::move(fields)), + _vars(std::move(vars)), + _bsonCurrent(bsonBegin) {} + +std::unique_ptr<PlanStage> BSONScanStage::clone() const { + return std::make_unique<BSONScanStage>(_bsonBegin, _bsonEnd, _recordSlot, _fields, _vars); +} + +void BSONScanStage::prepare(CompileCtx& ctx) { + if (_recordSlot) { + _recordAccessor = std::make_unique<value::ViewOfValueAccessor>(); + } + + for (size_t idx = 0; idx < _fields.size(); ++idx) { + auto [it, inserted] = + _fieldAccessors.emplace(_fields[idx], std::make_unique<value::ViewOfValueAccessor>()); + uassert(4822841, str::stream() << "duplicate field: " << _fields[idx], inserted); + auto [itRename, insertedRename] = _varAccessors.emplace(_vars[idx], it->second.get()); + uassert(4822842, str::stream() << "duplicate field: " << _vars[idx], insertedRename); + } +} + +value::SlotAccessor* BSONScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_recordSlot && *_recordSlot == slot) { + return _recordAccessor.get(); + } + + if (auto it = _varAccessors.find(slot); it != _varAccessors.end()) { + return it->second; + } + + return ctx.getAccessor(slot); +} + +void BSONScanStage::open(bool reOpen) { + _commonStats.opens++; + _bsonCurrent = _bsonBegin; +} + +PlanState BSONScanStage::getNext() { + if (_bsonCurrent < _bsonEnd) { + if (_recordAccessor) { + _recordAccessor->reset(value::TypeTags::bsonObject, + value::bitcastFrom<const char*>(_bsonCurrent)); + } + + if (auto fieldsToMatch = _fieldAccessors.size(); fieldsToMatch != 0) { + auto be = _bsonCurrent + 4; + auto end = _bsonCurrent + value::readFromMemory<uint32_t>(_bsonCurrent); + for (auto& [name, accessor] : _fieldAccessors) { + accessor->reset(); + } + while (*be != 0) { + auto sv = bson::fieldNameView(be); + if (auto it = _fieldAccessors.find(sv); it != _fieldAccessors.end()) { + // Found the field so convert it to Value. + auto [tag, val] = bson::convertFrom(true, be, end, sv.size()); + + it->second->reset(tag, val); + + if ((--fieldsToMatch) == 0) { + // No need to scan any further so bail out early. + break; + } + } + + be = bson::advance(be, sv.size()); + } + } + + // Advance to the next document. + _bsonCurrent += value::readFromMemory<uint32_t>(_bsonCurrent); + + _specificStats.numReads++; + return trackPlanState(PlanState::ADVANCED); + } + + _commonStats.isEOF = true; + return trackPlanState(PlanState::IS_EOF); +} + +void BSONScanStage::close() { + _commonStats.closes++; +} + +std::unique_ptr<PlanStageStats> BSONScanStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->specific = std::make_unique<ScanStats>(_specificStats); + return ret; +} + +const SpecificStats* BSONScanStage::getSpecificStats() const { + return &_specificStats; +} + +std::vector<DebugPrinter::Block> BSONScanStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + DebugPrinter::addKeyword(ret, "bsonscan"); + + + if (_recordSlot) { + DebugPrinter::addIdentifier(ret, _recordSlot.get()); + } + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _fields.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vars[idx]); + ret.emplace_back("="); + DebugPrinter::addIdentifier(ret, _fields[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + return ret; +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/bson_scan.h b/src/mongo/db/exec/sbe/stages/bson_scan.h new file mode 100644 index 00000000000..08a6c0aed77 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/bson_scan.h @@ -0,0 +1,76 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/values/bson.h" + +namespace mongo { +namespace sbe { +class BSONScanStage final : public PlanStage { +public: + BSONScanStage(const char* bsonBegin, + const char* bsonEnd, + boost::optional<value::SlotId> recordSlot, + std::vector<std::string> fields, + value::SlotVector vars); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + const char* const _bsonBegin; + const char* const _bsonEnd; + + const boost::optional<value::SlotId> _recordSlot; + const std::vector<std::string> _fields; + const value::SlotVector _vars; + + std::unique_ptr<value::ViewOfValueAccessor> _recordAccessor; + + value::FieldAccessorMap _fieldAccessors; + value::SlotAccessorMap _varAccessors; + + const char* _bsonCurrent; + + ScanStats _specificStats; +}; +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/check_bounds.cpp b/src/mongo/db/exec/sbe/stages/check_bounds.cpp new file mode 100644 index 00000000000..2caac57d879 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/check_bounds.cpp @@ -0,0 +1,146 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/check_bounds.h" + +namespace mongo::sbe { +CheckBoundsStage::CheckBoundsStage(std::unique_ptr<PlanStage> input, + const CheckBoundsParams& params, + value::SlotId inKeySlot, + value::SlotId inRecordIdSlot, + value::SlotId outSlot) + : PlanStage{"chkbounds"_sd}, + _params{params}, + _checker{&_params.bounds, _params.keyPattern, _params.direction}, + _inKeySlot{inKeySlot}, + _inRecordIdSlot{inRecordIdSlot}, + _outSlot{outSlot} { + _children.emplace_back(std::move(input)); +} + +std::unique_ptr<PlanStage> CheckBoundsStage::clone() const { + return std::make_unique<CheckBoundsStage>( + _children[0]->clone(), _params, _inKeySlot, _inRecordIdSlot, _outSlot); +} + +void CheckBoundsStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + _inKeyAccessor = _children[0]->getAccessor(ctx, _inKeySlot); + _inRecordIdAccessor = _children[0]->getAccessor(ctx, _inRecordIdSlot); +} + +value::SlotAccessor* CheckBoundsStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_outSlot == slot) { + return &_outAccessor; + } + + return _children[0]->getAccessor(ctx, slot); +} + +void CheckBoundsStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + _isEOF = false; +} + +PlanState CheckBoundsStage::getNext() { + if (_isEOF) { + return trackPlanState(PlanState::IS_EOF); + } + + auto state = _children[0]->getNext(); + + if (state == PlanState::ADVANCED) { + auto [keyTag, keyVal] = _inKeyAccessor->getViewOfValue(); + uassert(ErrorCodes::BadValue, "Wrong index key type", keyTag == value::TypeTags::ksValue); + + auto key = value::getKeyStringView(keyVal); + auto bsonKey = KeyString::toBson(*key, _params.ord); + IndexSeekPoint seekPoint; + + switch (_checker.checkKey(bsonKey, &seekPoint)) { + case IndexBoundsChecker::VALID: { + auto [tag, val] = _inRecordIdAccessor->getViewOfValue(); + _outAccessor.reset(false, tag, val); + break; + } + + case IndexBoundsChecker::DONE: + state = PlanState::IS_EOF; + break; + + case IndexBoundsChecker::MUST_ADVANCE: { + auto seekKey = std::make_unique<KeyString::Value>( + IndexEntryComparison::makeKeyStringFromSeekPointForSeek( + seekPoint, _params.version, _params.ord, _params.direction == 1)); + _outAccessor.reset( + true, value::TypeTags::ksValue, value::bitcastFrom(seekKey.release())); + // We should return the seek key provided by the 'IndexBoundsChecker' to restart + // the index scan, but we should stop on the next call to 'getNext' since we've just + // passed behind the current interval and need to signal the parent stage that we're + // done and can only continue further once the stage is reopened. + _isEOF = true; + break; + } + } + } + return trackPlanState(state); +} + +void CheckBoundsStage::close() { + _commonStats.closes++; + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> CheckBoundsStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* CheckBoundsStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> CheckBoundsStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "chkbounds"); + + DebugPrinter::addIdentifier(ret, _inKeySlot); + DebugPrinter::addIdentifier(ret, _inRecordIdSlot); + DebugPrinter::addIdentifier(ret, _outSlot); + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/check_bounds.h b/src/mongo/db/exec/sbe/stages/check_bounds.h new file mode 100644 index 00000000000..9d16f4f3507 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/check_bounds.h @@ -0,0 +1,100 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/query/index_bounds.h" + +namespace mongo::sbe { +struct CheckBoundsParams { + const IndexBounds bounds; + const BSONObj keyPattern; + const int direction; + const KeyString::Version version; + const Ordering ord; +}; + +/** + * This PlanStage takes a pair of index key/recordId slots as an input from its child stage (usually + * an index scan), and uses the 'IndexBoundsChecker' to check if the index key is within the index + * bounds specified in the 'CheckBoundsParams'. + * + * The stage is used when index bounds cannot be specified as valid low and high keys, which can be + * fed into the 'IndexScanStage' stage. For example, when index bounds are specified as + * multi-interval bounds and we cannot decompose it into a number of single-interval bounds. + * + * For each input pair, the stage can produce the following output, bound to the 'outSlot': + * + * 1. The key is within the bounds - caller can use data associated with this key, so the + * 'outSlot' would contain the recordId value. + * 2. The key is past the bounds - no further keys will satisfy the bounds and the caller should + * stop, so an EOF would be returned. + * 3. The key is not within the bounds, but has not exceeded the maximum value. The index scan + * would need to advance to the index key provided by the 'IndexBoundsChecker' and returned in + * the 'outSlot', and restart the scan from that key. + * + * This stage is usually used along with the stack spool to recursively feed the index key produced + * in case #3 back to the index scan, + */ +class CheckBoundsStage final : public PlanStage { +public: + CheckBoundsStage(std::unique_ptr<PlanStage> input, + const CheckBoundsParams& params, + value::SlotId inKeySlot, + value::SlotId inRecordIdSlot, + value::SlotId outSlot); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + const CheckBoundsParams _params; + IndexBoundsChecker _checker; + + const value::SlotId _inKeySlot; + const value::SlotId _inRecordIdSlot; + const value::SlotId _outSlot; + + value::SlotAccessor* _inKeyAccessor{nullptr}; + value::SlotAccessor* _inRecordIdAccessor{nullptr}; + value::OwnedValueAccessor _outAccessor; + + bool _isEOF{false}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/co_scan.cpp b/src/mongo/db/exec/sbe/stages/co_scan.cpp new file mode 100644 index 00000000000..4a2a3fa1406 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/co_scan.cpp @@ -0,0 +1,76 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/co_scan.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" + +namespace mongo::sbe { +CoScanStage::CoScanStage() : PlanStage("coscan"_sd) {} +std::unique_ptr<PlanStage> CoScanStage::clone() const { + return std::make_unique<CoScanStage>(); +} +void CoScanStage::prepare(CompileCtx& ctx) {} +value::SlotAccessor* CoScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + return ctx.getAccessor(slot); +} + +void CoScanStage::open(bool reOpen) { + _commonStats.opens++; +} + +PlanState CoScanStage::getNext() { + checkForInterrupt(_opCtx); + + // Run forever. + _commonStats.advances++; + return PlanState::ADVANCED; +} + +std::unique_ptr<PlanStageStats> CoScanStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + return ret; +} + +const SpecificStats* CoScanStage::getSpecificStats() const { + return nullptr; +} + +void CoScanStage::close() { + _commonStats.closes++; +} + +std::vector<DebugPrinter::Block> CoScanStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "coscan"); + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/co_scan.h b/src/mongo/db/exec/sbe/stages/co_scan.h new file mode 100644 index 00000000000..f88f3762852 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/co_scan.h @@ -0,0 +1,59 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/stages.h" + +namespace mongo::sbe { +/** + * This is a CoScan PlanStage. It delivers an infinite stream of getNext() calls. Also, it does not + * define any slots; i.e. it does not produce any results. + * + * On its face value this does not seem to be very useful but it is handy when we have to construct + * a data stream when there is not any physical source (i.e. no collection to read from). + * Typical use cases are: inner side of Traverse, outer side of Nested Loops, constants, etc. + */ +class CoScanStage final : public PlanStage { +public: + CoScanStage(); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/exchange.cpp b/src/mongo/db/exec/sbe/stages/exchange.cpp new file mode 100644 index 00000000000..e644c162e59 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/exchange.cpp @@ -0,0 +1,580 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/exchange.h" + +#include "mongo/base/init.h" +#include "mongo/db/client.h" + +namespace mongo::sbe { +std::unique_ptr<ThreadPool> s_globalThreadPool; +MONGO_INITIALIZER(s_globalThreadPool)(InitializerContext* context) { + ThreadPool::Options options; + options.poolName = "parallel execution pool"; + options.threadNamePrefix = "ExchProd"; + options.minThreads = 0; + options.maxThreads = 128; + options.onCreateThread = [](const std::string& name) { Client::initThread(name); }; + s_globalThreadPool = std::make_unique<ThreadPool>(options); + s_globalThreadPool->startup(); + + return Status::OK(); +} + +ExchangePipe::ExchangePipe(size_t size) { + // All buffers start empty. + _fullCount = 0; + _emptyCount = size; + for (size_t i = 0; i < _emptyCount; ++i) { + _fullBuffers.emplace_back(nullptr); + _emptyBuffers.emplace_back(std::make_unique<ExchangeBuffer>()); + } + + // Add a sentinel. + _fullBuffers.emplace_back(nullptr); +} + +void ExchangePipe::close() { + stdx::unique_lock lock(_mutex); + + _closed = true; + + _cond.notify_all(); +} + +std::unique_ptr<ExchangeBuffer> ExchangePipe::getEmptyBuffer() { + stdx::unique_lock lock(_mutex); + + _cond.wait(lock, [this]() { return _closed || _emptyCount > 0; }); + + if (_closed) { + return nullptr; + } + + --_emptyCount; + + return std::move(_emptyBuffers[_emptyCount]); +} + +std::unique_ptr<ExchangeBuffer> ExchangePipe::getFullBuffer() { + stdx::unique_lock lock(_mutex); + + _cond.wait(lock, [this]() { return _closed || _fullCount != _fullPosition; }); + + if (_closed) { + return nullptr; + } + + auto pos = _fullPosition; + _fullPosition = (_fullPosition + 1) % _fullBuffers.size(); + + return std::move(_fullBuffers[pos]); +} + +void ExchangePipe::putEmptyBuffer(std::unique_ptr<ExchangeBuffer> b) { + stdx::unique_lock lock(_mutex); + + _emptyBuffers[_emptyCount] = std::move(b); + + ++_emptyCount; + + _cond.notify_all(); +} + +void ExchangePipe::putFullBuffer(std::unique_ptr<ExchangeBuffer> b) { + stdx::unique_lock lock(_mutex); + + _fullBuffers[_fullCount] = std::move(b); + + _fullCount = (_fullCount + 1) % _fullBuffers.size(); + + _cond.notify_all(); +} + +ExchangeState::ExchangeState(size_t numOfProducers, + value::SlotVector fields, + ExchangePolicy policy, + std::unique_ptr<EExpression> partition, + std::unique_ptr<EExpression> orderLess) + : _policy(policy), + _numOfProducers(numOfProducers), + _fields(std::move(fields)), + _partition(std::move(partition)), + _orderLess(std::move(orderLess)) {} + +ExchangePipe* ExchangeState::pipe(size_t consumerTid, size_t producerTid) { + return _consumers[consumerTid]->pipe(producerTid); +} + +ExchangeBuffer* ExchangeConsumer::getBuffer(size_t producerId) { + if (_fullBuffers[producerId]) { + return _fullBuffers[producerId].get(); + } + + _fullBuffers[producerId] = _pipes[producerId]->getFullBuffer(); + + return _fullBuffers[producerId].get(); +} + +void ExchangeConsumer::putBuffer(size_t producerId) { + if (!_fullBuffers[producerId]) { + uasserted(4822832, "get not called before put"); + } + + // Clear the buffer before putting it back on the empty (free) list. + _fullBuffers[producerId]->clear(); + + _pipes[producerId]->putEmptyBuffer(std::move(_fullBuffers[producerId])); +} + +ExchangeConsumer::ExchangeConsumer(std::unique_ptr<PlanStage> input, + size_t numOfProducers, + value::SlotVector fields, + ExchangePolicy policy, + std::unique_ptr<EExpression> partition, + std::unique_ptr<EExpression> orderLess) + : PlanStage("exchange"_sd) { + _children.emplace_back(std::move(input)); + _state = std::make_shared<ExchangeState>( + numOfProducers, std::move(fields), policy, std::move(partition), std::move(orderLess)); + + _tid = _state->addConsumer(this); + _orderPreserving = _state->isOrderPreserving(); +} +ExchangeConsumer::ExchangeConsumer(std::shared_ptr<ExchangeState> state) + : PlanStage("exchange"_sd), _state(state) { + _tid = _state->addConsumer(this); + _orderPreserving = _state->isOrderPreserving(); +} +std::unique_ptr<PlanStage> ExchangeConsumer::clone() const { + return std::make_unique<ExchangeConsumer>(_state); +} +void ExchangeConsumer::prepare(CompileCtx& ctx) { + for (size_t idx = 0; idx < _state->fields().size(); ++idx) { + _outgoing.emplace_back(ExchangeBuffer::Accessor{}); + } + // Compile '<' function once we implement order preserving exchange. +} +value::SlotAccessor* ExchangeConsumer::getAccessor(CompileCtx& ctx, value::SlotId slot) { + // Accessors to pipes. + for (size_t idx = 0; idx < _state->fields().size(); ++idx) { + if (_state->fields()[idx] == slot) { + return &_outgoing[idx]; + } + } + + return ctx.getAccessor(slot); +} +void ExchangeConsumer::open(bool reOpen) { + _commonStats.opens++; + + if (reOpen) { + uasserted(4822833, "exchange consumer cannot be reopened"); + } + + { + stdx::unique_lock lock(_state->consumerOpenMutex()); + bool allConsumers = (++_state->consumerOpen()) == _state->numOfConsumers(); + + // Create all pipes. + if (_orderPreserving) { + for (size_t idx = 0; idx < _state->numOfProducers(); ++idx) { + _pipes.emplace_back(std::make_unique<ExchangePipe>(2)); + _fullBuffers.emplace_back(nullptr); + _bufferPos.emplace_back(0); + } + } else { + _pipes.emplace_back(std::make_unique<ExchangePipe>(_state->numOfProducers() * 2)); + _fullBuffers.emplace_back(nullptr); + _bufferPos.emplace_back(0); + } + _eofs = 0; + + if (_tid == 0) { + // Consumer ID 0 + + // Wait for all other consumers to show up. + if (!allConsumers) { + _state->consumerOpenCond().wait( + lock, [this]() { return _state->consumerOpen() == _state->numOfConsumers(); }); + } + + // Clone n copies of the subtree for every producer. + + PlanStage* masterSubTree = _children[0].get(); + masterSubTree->detachFromOperationContext(); + + for (size_t idx = 0; idx < _state->numOfProducers(); ++idx) { + if (idx == 0) { + _state->producerPlans().emplace_back( + std::make_unique<ExchangeProducer>(std::move(_children[0]), _state)); + // We have moved the child to the producer so clear the children vector. + _children.clear(); + } else { + _state->producerPlans().emplace_back( + std::make_unique<ExchangeProducer>(masterSubTree->clone(), _state)); + } + } + + // Start n producers. + for (size_t idx = 0; idx < _state->numOfProducers(); ++idx) { + auto pf = makePromiseFuture<void>(); + s_globalThreadPool->schedule( + [this, idx, promise = std::move(pf.promise)](auto status) mutable { + invariant(status); + + auto opCtx = cc().makeOperationContext(); + + promise.setWith([&] { + ExchangeProducer::start(opCtx.get(), + std::move(_state->producerPlans()[idx])); + }); + }); + _state->addProducerFuture(std::move(pf.future)); + } + } else { + // Consumer ID >0 + + // Make consumer 0 know that this consumer has shown up. + if (allConsumers) { + _state->consumerOpenCond().notify_all(); + } + } + } + + { + stdx::unique_lock lock(_state->consumerOpenMutex()); + if (_tid == 0) { + // Signal all other consumers that the open is done. + _state->consumerOpen() = 0; + _state->consumerOpenCond().notify_all(); + } else { + // Wait for the open to be done. + _state->consumerOpenCond().wait(lock, [this]() { return _state->consumerOpen() == 0; }); + } + } +} + +PlanState ExchangeConsumer::getNext() { + if (_orderPreserving) { + // Build a heap and return min element. + uasserted(4822834, "ordere exchange not yet implemented"); + } else { + while (_eofs < _state->numOfProducers()) { + auto buffer = getBuffer(0); + if (!buffer) { + // early out + return trackPlanState(PlanState::IS_EOF); + } + if (_bufferPos[0] < buffer->count()) { + // We still return from the current buffer. + for (size_t idx = 0; idx < _outgoing.size(); ++idx) { + _outgoing[idx].setBuffer(buffer); + _outgoing[idx].setIndex(_bufferPos[0] * _state->fields().size() + idx); + } + ++_bufferPos[0]; + ++_rowProcessed; + return trackPlanState(PlanState::ADVANCED); + } + + if (buffer->isEof()) { + ++_eofs; + } + + putBuffer(0); + _bufferPos[0] = 0; + } + } + return trackPlanState(PlanState::IS_EOF); +} +void ExchangeConsumer::close() { + _commonStats.closes++; + + { + stdx::unique_lock lock(_state->consumerCloseMutex()); + ++_state->consumerClose(); + + // Signal early out. + for (auto& p : _pipes) { + p->close(); + } + + if (_tid == 0) { + // Consumer ID 0 + // Wait for n producers to finish. + for (size_t idx = 0; idx < _state->numOfProducers(); ++idx) { + _state->producerResults()[idx].wait(); + } + } + + if (_state->consumerClose() == _state->numOfConsumers()) { + // Signal all other consumers that the close is done. + _state->consumerCloseCond().notify_all(); + } else { + // Wait for the close to be done. + _state->consumerCloseCond().wait( + lock, [this]() { return _state->consumerClose() == _state->numOfConsumers(); }); + } + } + // Rethrow the first stored exception from producers. + // We can do it outside of the lock as everybody else is gone by now. + if (_tid == 0) { + // Consumer ID 0 + for (size_t idx = 0; idx < _state->numOfProducers(); ++idx) { + _state->producerResults()[idx].get(); + } + } +} + +std::unique_ptr<PlanStageStats> ExchangeConsumer::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* ExchangeConsumer::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> ExchangeConsumer::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "exchange"); + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _state->fields().size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _state->fields()[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(std::to_string(_state->numOfProducers())); + + switch (_state->policy()) { + case ExchangePolicy::broadcast: + DebugPrinter::addKeyword(ret, "bcast"); + break; + case ExchangePolicy::roundrobin: + DebugPrinter::addKeyword(ret, "round"); + break; + default: + uasserted(4822835, "policy not yet implemented"); + } + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; +} + +ExchangePipe* ExchangeConsumer::pipe(size_t producerTid) { + if (_orderPreserving) { + return _pipes[producerTid].get(); + } else { + return _pipes[0].get(); + } +} + +ExchangeBuffer* ExchangeProducer::getBuffer(size_t consumerId) { + if (_emptyBuffers[consumerId]) { + return _emptyBuffers[consumerId].get(); + } + + _emptyBuffers[consumerId] = _pipes[consumerId]->getEmptyBuffer(); + + if (!_emptyBuffers[consumerId]) { + closePipes(); + } + + return _emptyBuffers[consumerId].get(); +} + +void ExchangeProducer::putBuffer(size_t consumerId) { + if (!_emptyBuffers[consumerId]) { + uasserted(4822836, "get not called before put"); + } + + _pipes[consumerId]->putFullBuffer(std::move(_emptyBuffers[consumerId])); +} + +void ExchangeProducer::closePipes() { + for (auto& p : _pipes) { + p->close(); + } +} + +ExchangeProducer::ExchangeProducer(std::unique_ptr<PlanStage> input, + std::shared_ptr<ExchangeState> state) + : PlanStage("exchangep"_sd), _state(state) { + _children.emplace_back(std::move(input)); + + _tid = _state->addProducer(this); + + // Retrieve the correct pipes. + for (size_t idx = 0; idx < _state->numOfConsumers(); ++idx) { + _pipes.emplace_back(_state->pipe(idx, _tid)); + _emptyBuffers.emplace_back(nullptr); + } +} + +void ExchangeProducer::start(OperationContext* opCtx, std::unique_ptr<PlanStage> producer) { + ExchangeProducer* p = static_cast<ExchangeProducer*>(producer.get()); + + p->attachFromOperationContext(opCtx); + + try { + CompileCtx ctx; + p->prepare(ctx); + p->open(false); + + auto status = p->getNext(); + if (status != PlanState::IS_EOF) { + uasserted(4822837, "producer returned invalid state"); + } + + p->close(); + } catch (...) { + // This is a bit sketchy but close the pipes as minimum. + p->closePipes(); + throw; + } +} + +std::unique_ptr<PlanStage> ExchangeProducer::clone() const { + uasserted(4822838, "ExchangeProducer is not cloneable"); +} + +void ExchangeProducer::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + for (auto& f : _state->fields()) { + _incoming.emplace_back(_children[0]->getAccessor(ctx, f)); + } +} +value::SlotAccessor* ExchangeProducer::getAccessor(CompileCtx& ctx, value::SlotId slot) { + return _children[0]->getAccessor(ctx, slot); +} +void ExchangeProducer::open(bool reOpen) { + _commonStats.opens++; + if (reOpen) { + uasserted(4822839, "exchange producer cannot be reopened"); + } + _children[0]->open(reOpen); +} +bool ExchangeProducer::appendData(size_t consumerId) { + auto buffer = getBuffer(consumerId); + // Detect early out. + if (!buffer) { + return false; + } + + // Copy data to buffer. + if (buffer->appendData(_incoming)) { + // Send it off to consumer when full. + putBuffer(consumerId); + } + + return true; +} + +PlanState ExchangeProducer::getNext() { + while (_children[0]->getNext() == PlanState::ADVANCED) { + // Push to the correct pipe. + switch (_state->policy()) { + case ExchangePolicy::broadcast: { + for (size_t idx = 0; idx < _pipes.size(); ++idx) { + // Detect early out in the loop. + if (!appendData(idx)) { + return trackPlanState(PlanState::IS_EOF); + } + } + } break; + case ExchangePolicy::roundrobin: { + // Detect early out. + if (!appendData(_roundRobinCounter)) { + return trackPlanState(PlanState::IS_EOF); + } + _roundRobinCounter = (_roundRobinCounter + 1) % _pipes.size(); + } break; + case ExchangePolicy::partition: { + uasserted(4822840, "policy not yet implemented"); + } break; + default: + MONGO_UNREACHABLE; + break; + } + } + + // Send off partially filled buffers and the eof marker. + for (size_t idx = 0; idx < _pipes.size(); ++idx) { + auto buffer = getBuffer(idx); + // Detect early out in the loop. + if (!buffer) { + return trackPlanState(PlanState::IS_EOF); + } + buffer->markEof(); + // Send it off to consumer. + putBuffer(idx); + } + return trackPlanState(PlanState::IS_EOF); +} +void ExchangeProducer::close() { + _commonStats.closes++; + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> ExchangeProducer::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* ExchangeProducer::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> ExchangeProducer::debugPrint() const { + return std::vector<DebugPrinter::Block>(); +} +bool ExchangeBuffer::appendData(std::vector<value::SlotAccessor*>& data) { + ++_count; + for (auto accesor : data) { + auto [tag, val] = accesor->copyOrMoveValue(); + value::ValueGuard guard{tag, val}; + _typeTags.push_back(tag); + _values.push_back(val); + guard.reset(); + } + + // A simply heuristic for now. + return isFull(); +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/exchange.h b/src/mongo/db/exec/sbe/stages/exchange.h new file mode 100644 index 00000000000..3fb3a3b91c5 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/exchange.h @@ -0,0 +1,331 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <vector> + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/stdx/condition_variable.h" +#include "mongo/stdx/future.h" +#include "mongo/util/concurrency/thread_pool.h" +#include "mongo/util/future.h" + +namespace mongo::sbe { +class ExchangeConsumer; +class ExchangeProducer; + +enum class ExchangePolicy { broadcast, roundrobin, partition }; + +// A unit of exchange between a consumer and a producer +class ExchangeBuffer { +public: + ~ExchangeBuffer() { + clear(); + } + void clear() { + _eof = false; + _count = 0; + + for (size_t idx = 0; idx < _typeTags.size(); ++idx) { + value::releaseValue(_typeTags[idx], _values[idx]); + } + _typeTags.clear(); + _values.clear(); + } + void markEof() { + _eof = true; + } + auto isEof() const { + return _eof; + } + + bool appendData(std::vector<value::SlotAccessor*>& data); + auto count() const { + return _count; + } + // The numbers are arbitrary for now. + auto isFull() const { + return _typeTags.size() >= 10240 || _count >= 1024; + } + + class Accessor : public value::SlotAccessor { + public: + void setBuffer(ExchangeBuffer* buffer) { + _buffer = buffer; + } + void setIndex(size_t index) { + _index = index; + } + + // Return non-owning view of the value. + std::pair<value::TypeTags, value::Value> getViewOfValue() const final { + return {_buffer->_typeTags[_index], _buffer->_values[_index]}; + } + std::pair<value::TypeTags, value::Value> copyOrMoveValue() final { + auto tag = _buffer->_typeTags[_index]; + auto val = _buffer->_values[_index]; + + _buffer->_typeTags[_index] = value::TypeTags::Nothing; + + return {tag, val}; + } + + private: + ExchangeBuffer* _buffer{nullptr}; + size_t _index{0}; + }; + +private: + std::vector<value::TypeTags> _typeTags; + std::vector<value::Value> _values; + + // Mark that this is the last buffer. + bool _eof{false}; + size_t _count{0}; + + friend class Accessor; +}; + +/** + * A connection that moves data between a consumer and a producer. + */ +class ExchangePipe { +public: + ExchangePipe(size_t size); + + void close(); + std::unique_ptr<ExchangeBuffer> getEmptyBuffer(); + std::unique_ptr<ExchangeBuffer> getFullBuffer(); + void putEmptyBuffer(std::unique_ptr<ExchangeBuffer>); + void putFullBuffer(std::unique_ptr<ExchangeBuffer>); + +private: + Mutex _mutex = MONGO_MAKE_LATCH("ExchangePipe::_mutex"); + stdx::condition_variable _cond; + + std::vector<std::unique_ptr<ExchangeBuffer>> _fullBuffers; + std::vector<std::unique_ptr<ExchangeBuffer>> _emptyBuffers; + size_t _fullCount{0}; + size_t _fullPosition{0}; + size_t _emptyCount{0}; + + // early out - pipe closed + bool _closed{false}; +}; + +/** + * Common shared state between all consumers and producers of a single exchange. + */ +class ExchangeState { +public: + ExchangeState(size_t numOfProducers, + value::SlotVector fields, + ExchangePolicy policy, + std::unique_ptr<EExpression> partition, + std::unique_ptr<EExpression> orderLess); + + bool isOrderPreserving() const { + return !!_orderLess; + } + auto policy() const { + return _policy; + } + + size_t addConsumer(ExchangeConsumer* c) { + _consumers.push_back(c); + return _consumers.size() - 1; + } + + size_t addProducer(ExchangeProducer* p) { + _producers.push_back(p); + return _producers.size() - 1; + } + + void addProducerFuture(Future<void> f) { + _producerResults.emplace_back(std::move(f)); + } + + auto& consumerOpenMutex() { + return _consumerOpenMutex; + } + auto& consumerOpenCond() { + return _consumerOpenCond; + } + auto& consumerOpen() { + return _consumerOpen; + } + + auto& consumerCloseMutex() { + return _consumerCloseMutex; + } + auto& consumerCloseCond() { + return _consumerCloseCond; + } + auto& consumerClose() { + return _consumerClose; + } + + auto& producerPlans() { + return _producerPlans; + } + auto& producerResults() { + return _producerResults; + } + + auto numOfConsumers() const { + return _consumers.size(); + } + auto numOfProducers() const { + return _numOfProducers; + } + + auto& fields() const { + return _fields; + } + ExchangePipe* pipe(size_t consumerTid, size_t producerTid); + +private: + const ExchangePolicy _policy; + const size_t _numOfProducers; + std::vector<ExchangeConsumer*> _consumers; + std::vector<ExchangeProducer*> _producers; + std::vector<std::unique_ptr<PlanStage>> _producerPlans; + std::vector<Future<void>> _producerResults; + + // Variables (fields) that pass through the exchange. + const value::SlotVector _fields; + + // Partitioning function. + const std::unique_ptr<EExpression> _partition; + + // The '<' function for order preserving exchange. + const std::unique_ptr<EExpression> _orderLess; + + // This is verbose and heavyweight. Recondsider something lighter + // at minimum try to share a single mutex (i.e. _stateMutex) if safe + mongo::Mutex _consumerOpenMutex; + stdx::condition_variable _consumerOpenCond; + size_t _consumerOpen{0}; + + mongo::Mutex _consumerCloseMutex; + stdx::condition_variable _consumerCloseCond; + size_t _consumerClose{0}; +}; + +class ExchangeConsumer final : public PlanStage { +public: + ExchangeConsumer(std::unique_ptr<PlanStage> input, + size_t numOfProducers, + value::SlotVector fields, + ExchangePolicy policy, + std::unique_ptr<EExpression> partition, + std::unique_ptr<EExpression> orderLess); + + ExchangeConsumer(std::shared_ptr<ExchangeState> state); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + + ExchangePipe* pipe(size_t producerTid); + +private: + ExchangeBuffer* getBuffer(size_t producerId); + void putBuffer(size_t producerId); + + std::shared_ptr<ExchangeState> _state; + size_t _tid{0}; + + // Accessors for the outgoing values (from the exchange buffers). + std::vector<ExchangeBuffer::Accessor> _outgoing; + + // Pipes to producers (if order preserving) or a single pipe shared by all producers. + std::vector<std::unique_ptr<ExchangePipe>> _pipes; + + // Current full buffers that this consumer is processing. + std::vector<std::unique_ptr<ExchangeBuffer>> _fullBuffers; + + // Current position in buffers that this consumer is processing. + std::vector<size_t> _bufferPos; + + // Count how may EOFs we have seen so far. + size_t _eofs{0}; + + bool _orderPreserving{false}; + + size_t _rowProcessed{0}; +}; + +class ExchangeProducer final : public PlanStage { +public: + ExchangeProducer(std::unique_ptr<PlanStage> input, std::shared_ptr<ExchangeState> state); + + static void start(OperationContext* opCtx, std::unique_ptr<PlanStage> producer); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + ExchangeBuffer* getBuffer(size_t consumerId); + void putBuffer(size_t consumerId); + + void closePipes(); + bool appendData(size_t consumerId); + + std::shared_ptr<ExchangeState> _state; + size_t _tid{0}; + size_t _roundRobinCounter{0}; + + std::vector<value::SlotAccessor*> _incoming; + + std::vector<ExchangePipe*> _pipes; + + // Current empty buffers that this producer is processing. + std::vector<std::unique_ptr<ExchangeBuffer>> _emptyBuffers; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/filter.h b/src/mongo/db/exec/sbe/stages/filter.h new file mode 100644 index 00000000000..c8c87b5a34b --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/filter.h @@ -0,0 +1,167 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/vm/vm.h" + +namespace mongo::sbe { +/** + * This is a filter plan stage. If the IsConst template parameter is true then the filter expression + * is 'constant' i.e. it does not depend on values coming from its input. It means that we can + * evaluate it in the open() call and skip getNext() calls completely if the result is false. + * The IsEof template parameter controls 'early out' behavior of the filter expression. Once the + * filter evaluates to false then the getNext() call returns EOF. + */ +template <bool IsConst, bool IsEof = false> +class FilterStage final : public PlanStage { +public: + FilterStage(std::unique_ptr<PlanStage> input, std::unique_ptr<EExpression> filter) + : PlanStage(IsConst ? "cfilter"_sd : (IsEof ? "efilter" : "filter"_sd)), + _filter(std::move(filter)) { + static_assert(!IsEof || !IsConst); + _children.emplace_back(std::move(input)); + } + + std::unique_ptr<PlanStage> clone() const final { + return std::make_unique<FilterStage>(_children[0]->clone(), _filter->clone()); + } + + void prepare(CompileCtx& ctx) final { + _children[0]->prepare(ctx); + + ctx.root = this; + _filterCode = _filter->compile(ctx); + } + + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final { + return _children[0]->getAccessor(ctx, slot); + } + + void open(bool reOpen) final { + _commonStats.opens++; + + if constexpr (IsConst) { + _specificStats.numTested++; + + auto pass = _bytecode.runPredicate(_filterCode.get()); + if (!pass) { + close(); + return; + } + } + _children[0]->open(reOpen); + _childOpened = true; + } + + PlanState getNext() final { + // The constant filter evaluates the predicate in the open method. + if constexpr (IsConst) { + if (!_childOpened) { + return trackPlanState(PlanState::IS_EOF); + } else { + return trackPlanState(_children[0]->getNext()); + } + } + + auto state = PlanState::IS_EOF; + bool pass = false; + + do { + state = _children[0]->getNext(); + + if (state == PlanState::ADVANCED) { + _specificStats.numTested++; + + pass = _bytecode.runPredicate(_filterCode.get()); + + if constexpr (IsEof) { + if (!pass) { + return trackPlanState(PlanState::IS_EOF); + } + } + } + } while (state == PlanState::ADVANCED && !pass); + + return trackPlanState(state); + } + + void close() final { + _commonStats.closes++; + + if (_childOpened) { + _children[0]->close(); + _childOpened = false; + } + } + + std::unique_ptr<PlanStageStats> getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->specific = std::make_unique<FilterStats>(_specificStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; + } + + const SpecificStats* getSpecificStats() const final { + return &_specificStats; + } + + std::vector<DebugPrinter::Block> debugPrint() const final { + std::vector<DebugPrinter::Block> ret; + if constexpr (IsConst) { + DebugPrinter::addKeyword(ret, "cfilter"); + } else if constexpr (IsEof) { + DebugPrinter::addKeyword(ret, "efilter"); + } else { + DebugPrinter::addKeyword(ret, "filter"); + } + + ret.emplace_back("{`"); + DebugPrinter::addBlocks(ret, _filter->debugPrint()); + ret.emplace_back("`}"); + + DebugPrinter::addNewLine(ret); + + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; + } + +private: + const std::unique_ptr<EExpression> _filter; + std::unique_ptr<vm::CodeFragment> _filterCode; + + vm::ByteCode _bytecode; + + bool _childOpened{false}; + FilterStats _specificStats; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp new file mode 100644 index 00000000000..e8cfc33fde2 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp @@ -0,0 +1,199 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/hash_agg.h" + +#include "mongo/util/str.h" + +namespace mongo { +namespace sbe { +HashAggStage::HashAggStage(std::unique_ptr<PlanStage> input, + value::SlotVector gbs, + value::SlotMap<std::unique_ptr<EExpression>> aggs) + : PlanStage("group"_sd), _gbs(std::move(gbs)), _aggs(std::move(aggs)) { + _children.emplace_back(std::move(input)); +} + +std::unique_ptr<PlanStage> HashAggStage::clone() const { + value::SlotMap<std::unique_ptr<EExpression>> aggs; + for (auto& [k, v] : _aggs) { + aggs.emplace(k, v->clone()); + } + return std::make_unique<HashAggStage>(_children[0]->clone(), _gbs, std::move(aggs)); +} + +void HashAggStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + value::SlotSet dupCheck; + size_t counter = 0; + // Process group by columns. + for (auto& slot : _gbs) { + auto [it, inserted] = dupCheck.emplace(slot); + uassert(4822827, str::stream() << "duplicate field: " << slot, inserted); + + _inKeyAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + _outKeyAccessors.emplace_back(std::make_unique<HashKeyAccessor>(_htIt, counter++)); + _outAccessors[slot] = _outKeyAccessors.back().get(); + } + + counter = 0; + for (auto& [slot, expr] : _aggs) { + auto [it, inserted] = dupCheck.emplace(slot); + // Some compilers do not allow to capture local bindings by lambda functions (the one + // is used implicitly in uassert below), so we need a local variable to construct an + // error message. + const auto slotId = slot; + uassert(4822828, str::stream() << "duplicate field: " << slotId, inserted); + + _outAggAccessors.emplace_back(std::make_unique<HashAggAccessor>(_htIt, counter++)); + _outAccessors[slot] = _outAggAccessors.back().get(); + + ctx.root = this; + ctx.aggExpression = true; + ctx.accumulator = _outAggAccessors.back().get(); + + _aggCodes.emplace_back(expr->compile(ctx)); + ctx.aggExpression = false; + } + _compiled = true; +} + +value::SlotAccessor* HashAggStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_compiled) { + if (auto it = _outAccessors.find(slot); it != _outAccessors.end()) { + return it->second; + } + } else { + return _children[0]->getAccessor(ctx, slot); + } + + return ctx.getAccessor(slot); +} + +void HashAggStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + + value::MaterializedRow key; + while (_children[0]->getNext() == PlanState::ADVANCED) { + key._fields.resize(_inKeyAccessors.size()); + // Copy keys in order to do the lookup. + size_t idx = 0; + for (auto& p : _inKeyAccessors) { + auto [tag, val] = p->getViewOfValue(); + key._fields[idx++].reset(false, tag, val); + } + + auto [it, inserted] = _ht.emplace(std::move(key), value::MaterializedRow{}); + if (inserted) { + // Copy keys. + const_cast<value::MaterializedRow&>(it->first).makeOwned(); + // Initialize accumulators. + it->second._fields.resize(_outAggAccessors.size()); + } + + // Accumulate. + _htIt = it; + for (size_t idx = 0; idx < _outAggAccessors.size(); ++idx) { + auto [owned, tag, val] = _bytecode.run(_aggCodes[idx].get()); + _outAggAccessors[idx]->reset(owned, tag, val); + } + } + + _children[0]->close(); + + _htIt = _ht.end(); +} + +PlanState HashAggStage::getNext() { + if (_htIt == _ht.end()) { + _htIt = _ht.begin(); + } else { + ++_htIt; + } + + if (_htIt == _ht.end()) { + return trackPlanState(PlanState::IS_EOF); + } + + return trackPlanState(PlanState::ADVANCED); +} + +std::unique_ptr<PlanStageStats> HashAggStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* HashAggStage::getSpecificStats() const { + return nullptr; +} + +void HashAggStage::close() { + _commonStats.closes++; +} + +std::vector<DebugPrinter::Block> HashAggStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "group"); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _gbs.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _gbs[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block("[`")); + bool first = true; + for (auto& p : _aggs) { + if (!first) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, p.first); + ret.emplace_back("="); + DebugPrinter::addBlocks(ret, p.second->debugPrint()); + first = false; + } + ret.emplace_back("`]"); + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.h b/src/mongo/db/exec/sbe/stages/hash_agg.h new file mode 100644 index 00000000000..d36cdb13115 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/hash_agg.h @@ -0,0 +1,84 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <unordered_map> + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/vm/vm.h" +#include "mongo/stdx/unordered_map.h" + +namespace mongo { +namespace sbe { +class HashAggStage final : public PlanStage { +public: + HashAggStage(std::unique_ptr<PlanStage> input, + value::SlotVector gbs, + value::SlotMap<std::unique_ptr<EExpression>> aggs); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + using TableType = stdx:: + unordered_map<value::MaterializedRow, value::MaterializedRow, value::MaterializedRowHasher>; + + using HashKeyAccessor = value::MaterializedRowKeyAccessor<TableType::iterator>; + using HashAggAccessor = value::MaterializedRowValueAccessor<TableType::iterator>; + + const value::SlotVector _gbs; + const value::SlotMap<std::unique_ptr<EExpression>> _aggs; + + value::SlotAccessorMap _outAccessors; + std::vector<value::SlotAccessor*> _inKeyAccessors; + std::vector<std::unique_ptr<HashKeyAccessor>> _outKeyAccessors; + + std::vector<std::unique_ptr<HashAggAccessor>> _outAggAccessors; + std::vector<std::unique_ptr<vm::CodeFragment>> _aggCodes; + + TableType _ht; + TableType::iterator _htIt; + + vm::ByteCode _bytecode; + + bool _compiled{false}; +}; +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/hash_join.cpp b/src/mongo/db/exec/sbe/stages/hash_join.cpp new file mode 100644 index 00000000000..69e1dffe1e6 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/hash_join.cpp @@ -0,0 +1,263 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/hash_join.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/util/str.h" + +namespace mongo { +namespace sbe { +HashJoinStage::HashJoinStage(std::unique_ptr<PlanStage> outer, + std::unique_ptr<PlanStage> inner, + value::SlotVector outerCond, + value::SlotVector outerProjects, + value::SlotVector innerCond, + value::SlotVector innerProjects) + : PlanStage("hj"_sd), + _outerCond(std::move(outerCond)), + _outerProjects(std::move(outerProjects)), + _innerCond(std::move(innerCond)), + _innerProjects(std::move(innerProjects)) { + if (_outerCond.size() != _innerCond.size()) { + uasserted(4822823, "left and right size do not match"); + } + + _children.emplace_back(std::move(outer)); + _children.emplace_back(std::move(inner)); +} + +std::unique_ptr<PlanStage> HashJoinStage::clone() const { + return std::make_unique<HashJoinStage>(_children[0]->clone(), + _children[1]->clone(), + _outerCond, + _outerProjects, + _innerCond, + _innerProjects); +} + +void HashJoinStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + _children[1]->prepare(ctx); + + size_t counter = 0; + value::SlotSet dupCheck; + for (auto& slot : _outerCond) { + auto [it, inserted] = dupCheck.emplace(slot); + uassert(4822824, str::stream() << "duplicate field: " << slot, inserted); + + _inOuterKeyAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + _outOuterKeyAccessors.emplace_back(std::make_unique<HashKeyAccessor>(_htIt, counter++)); + _outOuterAccessors[slot] = _outOuterKeyAccessors.back().get(); + } + + counter = 0; + for (auto& slot : _innerCond) { + auto [it, inserted] = dupCheck.emplace(slot); + uassert(4822825, str::stream() << "duplicate field: " << slot, inserted); + + _inInnerKeyAccessors.emplace_back(_children[1]->getAccessor(ctx, slot)); + } + + counter = 0; + for (auto& slot : _outerProjects) { + auto [it, inserted] = dupCheck.emplace(slot); + uassert(4822826, str::stream() << "duplicate field: " << slot, inserted); + + _inOuterProjectAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + _outOuterProjectAccessors.emplace_back( + std::make_unique<HashProjectAccessor>(_htIt, counter++)); + _outOuterAccessors[slot] = _outOuterProjectAccessors.back().get(); + } + + _probeKey._fields.resize(_inInnerKeyAccessors.size()); + + _compiled = true; +} + +value::SlotAccessor* HashJoinStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_compiled) { + if (auto it = _outOuterAccessors.find(slot); it != _outOuterAccessors.end()) { + return it->second; + } + + return _children[1]->getAccessor(ctx, slot); + } + + return ctx.getAccessor(slot); +} + +void HashJoinStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + // Insert the outer side into the hash table. + value::MaterializedRow key; + value::MaterializedRow project; + + while (_children[0]->getNext() == PlanState::ADVANCED) { + key._fields.reserve(_inOuterKeyAccessors.size()); + project._fields.reserve(_inOuterProjectAccessors.size()); + + // Copy keys in order to do the lookup. + for (auto& p : _inOuterKeyAccessors) { + key._fields.push_back(value::OwnedValueAccessor{}); + auto [tag, val] = p->copyOrMoveValue(); + key._fields.back().reset(true, tag, val); + } + + // Copy projects. + for (auto& p : _inOuterProjectAccessors) { + project._fields.push_back(value::OwnedValueAccessor{}); + auto [tag, val] = p->copyOrMoveValue(); + project._fields.back().reset(true, tag, val); + } + + _ht.emplace(std::move(key), std::move(project)); + } + + _children[0]->close(); + + _children[1]->open(reOpen); + + _htIt = _ht.end(); + _htItEnd = _ht.end(); +} + +PlanState HashJoinStage::getNext() { + if (_htIt != _htItEnd) { + ++_htIt; + } + + if (_htIt == _htItEnd) { + while (_htIt == _htItEnd) { + auto state = _children[1]->getNext(); + if (state == PlanState::IS_EOF) { + // LEFT and OUTER joins should enumerate "non-returned" rows here. + return trackPlanState(state); + } + + // Copy keys in order to do the lookup. + size_t idx = 0; + for (auto& p : _inInnerKeyAccessors) { + auto [tag, val] = p->getViewOfValue(); + _probeKey._fields[idx++].reset(false, tag, val); + } + + auto [low, hi] = _ht.equal_range(_probeKey); + _htIt = low; + _htItEnd = hi; + // If _htIt == _htItEnd (i.e. no match) then RIGHT and OUTER joins + // should enumerate "non-returned" rows here. + } + } + + return trackPlanState(PlanState::ADVANCED); +} + +void HashJoinStage::close() { + _commonStats.closes++; + _children[1]->close(); +} + +std::unique_ptr<PlanStageStats> HashJoinStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + ret->children.emplace_back(_children[1]->getStats()); + return ret; +} + +const SpecificStats* HashJoinStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> HashJoinStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "hj"); + + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + + DebugPrinter::addKeyword(ret, "left"); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _outerCond.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _outerCond[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _outerProjects.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _outerProjects[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + DebugPrinter::addKeyword(ret, "right"); + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _innerCond.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _innerCond[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _innerProjects.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _innerProjects[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + DebugPrinter::addBlocks(ret, _children[1]->debugPrint()); + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + return ret; +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/hash_join.h b/src/mongo/db/exec/sbe/stages/hash_join.h new file mode 100644 index 00000000000..b71241cda3f --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/hash_join.h @@ -0,0 +1,101 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <vector> + +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/vm/vm.h" + +namespace mongo::sbe { +class HashJoinStage final : public PlanStage { +public: + HashJoinStage(std::unique_ptr<PlanStage> outer, + std::unique_ptr<PlanStage> inner, + value::SlotVector outerCond, + value::SlotVector outerProjects, + value::SlotVector innerCond, + value::SlotVector innerProjects); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + using TableType = std::unordered_multimap<value::MaterializedRow, // NOLINT + value::MaterializedRow, + value::MaterializedRowHasher>; + + using HashKeyAccessor = value::MaterializedRowKeyAccessor<TableType::iterator>; + using HashProjectAccessor = value::MaterializedRowValueAccessor<TableType::iterator>; + + const value::SlotVector _outerCond; + const value::SlotVector _outerProjects; + const value::SlotVector _innerCond; + const value::SlotVector _innerProjects; + + // All defined values from the outer side (i.e. they come from the hash table). + value::SlotAccessorMap _outOuterAccessors; + + // Accessors of input codition values (keys) that are being inserted into the hash table. + std::vector<value::SlotAccessor*> _inOuterKeyAccessors; + + // Accessors of output keys. + std::vector<std::unique_ptr<HashKeyAccessor>> _outOuterKeyAccessors; + + // Accessors of input projection values that are build inserted into the hash table. + std::vector<value::SlotAccessor*> _inOuterProjectAccessors; + + // Accessors of output projections. + std::vector<std::unique_ptr<HashProjectAccessor>> _outOuterProjectAccessors; + + // Accessors of input codition values (keys) that are being inserted into the hash table. + std::vector<value::SlotAccessor*> _inInnerKeyAccessors; + + // Key used to probe inside the hash table. + value::MaterializedRow _probeKey; + + TableType _ht; + TableType::iterator _htIt; + TableType::iterator _htItEnd; + + vm::ByteCode _bytecode; + + bool _compiled{false}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp new file mode 100644 index 00000000000..9e3cdb21d44 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp @@ -0,0 +1,334 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/ix_scan.h" + +#include "mongo/db/catalog/index_catalog.h" +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/index/index_access_method.h" + +namespace mongo::sbe { +IndexScanStage::IndexScanStage(const NamespaceStringOrUUID& name, + std::string_view indexName, + bool forward, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + boost::optional<value::SlotId> seekKeySlotLow, + boost::optional<value::SlotId> seekKeySlotHi, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) + : PlanStage(seekKeySlotLow ? "ixseek"_sd : "ixscan"_sd, yieldPolicy), + _name(name), + _indexName(indexName), + _forward(forward), + _recordSlot(recordSlot), + _recordIdSlot(recordIdSlot), + _fields(std::move(fields)), + _vars(std::move(vars)), + _seekKeySlotLow(seekKeySlotLow), + _seekKeySlotHi(seekKeySlotHi), + _tracker(tracker) { + invariant(_fields.size() == _vars.size()); + // The valid state is when both boundaries, or none is set, or only low key is set. + invariant((_seekKeySlotLow && _seekKeySlotHi) || (!_seekKeySlotLow && !_seekKeySlotHi) || + (_seekKeySlotLow && !_seekKeySlotHi)); +} + +std::unique_ptr<PlanStage> IndexScanStage::clone() const { + return std::make_unique<IndexScanStage>(_name, + _indexName, + _forward, + _recordSlot, + _recordIdSlot, + _fields, + _vars, + _seekKeySlotLow, + _seekKeySlotHi, + _yieldPolicy, + _tracker); +} + +void IndexScanStage::prepare(CompileCtx& ctx) { + if (_recordSlot) { + _recordAccessor = std::make_unique<value::ViewOfValueAccessor>(); + } + + if (_recordIdSlot) { + _recordIdAccessor = std::make_unique<value::ViewOfValueAccessor>(); + } + + for (size_t idx = 0; idx < _fields.size(); ++idx) { + auto [it, inserted] = + _fieldAccessors.emplace(_fields[idx], std::make_unique<value::ViewOfValueAccessor>()); + uassert(4822821, str::stream() << "duplicate field: " << _fields[idx], inserted); + auto [itRename, insertedRename] = _varAccessors.emplace(_vars[idx], it->second.get()); + uassert(4822822, str::stream() << "duplicate field: " << _vars[idx], insertedRename); + } + + if (_seekKeySlotLow) { + _seekKeyLowAccessor = ctx.getAccessor(*_seekKeySlotLow); + } + if (_seekKeySlotHi) { + _seekKeyHiAccessor = ctx.getAccessor(*_seekKeySlotHi); + } +} + +value::SlotAccessor* IndexScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_recordSlot && *_recordSlot == slot) { + return _recordAccessor.get(); + } + + if (_recordIdSlot && *_recordIdSlot == slot) { + return _recordIdAccessor.get(); + } + + if (auto it = _varAccessors.find(slot); it != _varAccessors.end()) { + return it->second; + } + + return ctx.getAccessor(slot); +} + +void IndexScanStage::doSaveState() { + if (_cursor) { + _cursor->save(); + } + + _coll.reset(); +} +void IndexScanStage::doRestoreState() { + invariant(_opCtx); + invariant(!_coll); + + // If this stage is not currently open, then there is nothing to restore. + if (!_open) { + return; + } + + _coll.emplace(_opCtx, _name); + + if (_cursor) { + _cursor->restore(); + } +} + +void IndexScanStage::doDetachFromOperationContext() { + if (_cursor) { + _cursor->detachFromOperationContext(); + } +} +void IndexScanStage::doAttachFromOperationContext(OperationContext* opCtx) { + if (_cursor) { + _cursor->reattachToOperationContext(opCtx); + } +} + +void IndexScanStage::open(bool reOpen) { + _commonStats.opens++; + + invariant(_opCtx); + if (!reOpen) { + invariant(!_cursor); + invariant(!_coll); + _coll.emplace(_opCtx, _name); + } else { + invariant(_cursor); + invariant(_coll); + } + + _open = true; + _firstGetNext = true; + + if (auto collection = _coll->getCollection()) { + auto indexCatalog = collection->getIndexCatalog(); + auto indexDesc = indexCatalog->findIndexByName(_opCtx, _indexName); + if (indexDesc) { + _weakIndexCatalogEntry = indexCatalog->getEntryShared(indexDesc); + } + + if (auto entry = _weakIndexCatalogEntry.lock()) { + if (!_cursor) { + _cursor = + entry->accessMethod()->getSortedDataInterface()->newCursor(_opCtx, _forward); + } + + if (_seekKeyLowAccessor && _seekKeyHiAccessor) { + auto [tagLow, valLow] = _seekKeyLowAccessor->getViewOfValue(); + uassert(4822851, "seek key is wrong type", tagLow == value::TypeTags::ksValue); + _seekKeyLow = value::getKeyStringView(valLow); + + auto [tagHi, valHi] = _seekKeyHiAccessor->getViewOfValue(); + uassert(4822852, "seek key is wrong type", tagHi == value::TypeTags::ksValue); + _seekKeyHi = value::getKeyStringView(valHi); + } else if (_seekKeyLowAccessor) { + auto [tagLow, valLow] = _seekKeyLowAccessor->getViewOfValue(); + uassert(4822853, "seek key is wrong type", tagLow == value::TypeTags::ksValue); + _seekKeyLow = value::getKeyStringView(valLow); + _seekKeyHi = nullptr; + } else { + auto sdi = entry->accessMethod()->getSortedDataInterface(); + KeyString::Builder kb(sdi->getKeyStringVersion(), + sdi->getOrdering(), + KeyString::Discriminator::kExclusiveBefore); + kb.appendDiscriminator(KeyString::Discriminator::kExclusiveBefore); + _startPoint = kb.getValueCopy(); + + _seekKeyLow = &_startPoint; + _seekKeyHi = nullptr; + } + } else { + _cursor.reset(); + } + } else { + _cursor.reset(); + } +} + +PlanState IndexScanStage::getNext() { + if (!_cursor) { + return trackPlanState(PlanState::IS_EOF); + } + + checkForInterrupt(_opCtx); + + if (_firstGetNext) { + _firstGetNext = false; + _nextRecord = _cursor->seekForKeyString(*_seekKeyLow); + } else { + _nextRecord = _cursor->nextKeyString(); + } + + if (!_nextRecord) { + return trackPlanState(PlanState::IS_EOF); + } + + if (_seekKeyHi) { + auto cmp = _nextRecord->keyString.compare(*_seekKeyHi); + + if (_forward) { + if (cmp > 0) { + return trackPlanState(PlanState::IS_EOF); + } + } else { + if (cmp < 0) { + return trackPlanState(PlanState::IS_EOF); + } + } + } + + if (_recordAccessor) { + _recordAccessor->reset(value::TypeTags::ksValue, + value::bitcastFrom(&_nextRecord->keyString)); + } + + if (_recordIdAccessor) { + _recordIdAccessor->reset(value::TypeTags::NumberInt64, + value::bitcastFrom<int64_t>(_nextRecord->loc.repr())); + } + + if (_tracker && _tracker->trackProgress<TrialRunProgressTracker::kNumReads>(1)) { + // If we're collecting execution stats during multi-planning and reached the end of the + // trial period (trackProgress() will return 'true' in this case), then we can reset the + // tracker. Note that a trial period is executed only once per a PlanStge tree, and once + // completed never run again on the same tree. + _tracker = nullptr; + } + ++_specificStats.numReads; + return trackPlanState(PlanState::ADVANCED); +} + +void IndexScanStage::close() { + _commonStats.closes++; + + _cursor.reset(); + _coll.reset(); + _open = false; +} + +std::unique_ptr<PlanStageStats> IndexScanStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->specific = std::make_unique<IndexScanStats>(_specificStats); + return ret; +} + +const SpecificStats* IndexScanStage::getSpecificStats() const { + return &_specificStats; +} + +std::vector<DebugPrinter::Block> IndexScanStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + if (_seekKeySlotLow) { + DebugPrinter::addKeyword(ret, "ixseek"); + + DebugPrinter::addIdentifier(ret, _seekKeySlotLow.get()); + if (_seekKeySlotHi) { + DebugPrinter::addIdentifier(ret, _seekKeySlotHi.get()); + } + } else { + DebugPrinter::addKeyword(ret, "ixscan"); + } + + if (_recordSlot) { + DebugPrinter::addIdentifier(ret, _recordSlot.get()); + } + + if (_recordIdSlot) { + DebugPrinter::addIdentifier(ret, _recordIdSlot.get()); + } + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _fields.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vars[idx]); + ret.emplace_back("="); + DebugPrinter::addIdentifier(ret, _fields[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back("@\"`"); + DebugPrinter::addIdentifier(ret, _name.toString()); + ret.emplace_back("`\""); + + ret.emplace_back("@\"`"); + DebugPrinter::addIdentifier(ret, _indexName); + ret.emplace_back("`\""); + + ret.emplace_back(_forward ? "true" : "false"); + + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h new file mode 100644 index 00000000000..a3d23b2bb82 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/ix_scan.h @@ -0,0 +1,108 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/db_raii.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/trial_run_progress_tracker.h" +#include "mongo/db/storage/record_store.h" +#include "mongo/db/storage/sorted_data_interface.h" + +namespace mongo::sbe { +class IndexScanStage final : public PlanStage { +public: + IndexScanStage(const NamespaceStringOrUUID& name, + std::string_view indexName, + bool forward, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + boost::optional<value::SlotId> seekKeySlotLow, + boost::optional<value::SlotId> seekKeySlotHi, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +protected: + void doSaveState() override; + void doRestoreState() override; + void doDetachFromOperationContext() override; + void doAttachFromOperationContext(OperationContext* opCtx) override; + +private: + const NamespaceStringOrUUID _name; + const std::string _indexName; + const bool _forward; + const boost::optional<value::SlotId> _recordSlot; + const boost::optional<value::SlotId> _recordIdSlot; + const std::vector<std::string> _fields; + const value::SlotVector _vars; + const boost::optional<value::SlotId> _seekKeySlotLow; + const boost::optional<value::SlotId> _seekKeySlotHi; + + std::unique_ptr<value::ViewOfValueAccessor> _recordAccessor; + std::unique_ptr<value::ViewOfValueAccessor> _recordIdAccessor; + + value::FieldAccessorMap _fieldAccessors; + value::SlotAccessorMap _varAccessors; + + value::SlotAccessor* _seekKeyLowAccessor{nullptr}; + value::SlotAccessor* _seekKeyHiAccessor{nullptr}; + + KeyString::Value _startPoint; + KeyString::Value* _seekKeyLow{nullptr}; + KeyString::Value* _seekKeyHi{nullptr}; + + std::unique_ptr<SortedDataInterface::Cursor> _cursor; + std::weak_ptr<const IndexCatalogEntry> _weakIndexCatalogEntry; + boost::optional<AutoGetCollectionForRead> _coll; + boost::optional<KeyStringEntry> _nextRecord; + + bool _open{false}; + bool _firstGetNext{true}; + IndexScanStats _specificStats; + + // If provided, used during a trial run to accumulate certain execution stats. Once the trial + // run is complete, this pointer is reset to nullptr. + TrialRunProgressTracker* _tracker{nullptr}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/limit_skip.cpp b/src/mongo/db/exec/sbe/stages/limit_skip.cpp new file mode 100644 index 00000000000..36cf2964631 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/limit_skip.cpp @@ -0,0 +1,113 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/limit_skip.h" + +namespace mongo::sbe { +LimitSkipStage::LimitSkipStage(std::unique_ptr<PlanStage> input, + boost::optional<long long> limit, + boost::optional<long long> skip) + : PlanStage(!skip ? "limit"_sd : "limitskip"_sd), + _limit(limit), + _skip(skip), + _current(0), + _isEOF(false) { + invariant(_limit || _skip); + _children.emplace_back(std::move(input)); + _specificStats.limit = limit; + _specificStats.skip = skip; +} + +std::unique_ptr<PlanStage> LimitSkipStage::clone() const { + return std::make_unique<LimitSkipStage>(_children[0]->clone(), _limit, _skip); +} + +void LimitSkipStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); +} + +value::SlotAccessor* LimitSkipStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + return _children[0]->getAccessor(ctx, slot); +} + +void LimitSkipStage::open(bool reOpen) { + _commonStats.opens++; + _isEOF = false; + _children[0]->open(reOpen); + + if (_skip) { + for (_current = 0; _current < *_skip && !_isEOF; _current++) { + _isEOF = _children[0]->getNext() == PlanState::IS_EOF; + } + } + _current = 0; +} +PlanState LimitSkipStage::getNext() { + if (_isEOF || (_limit && _current++ == *_limit)) { + return trackPlanState(PlanState::IS_EOF); + } + + return trackPlanState(_children[0]->getNext()); +} +void LimitSkipStage::close() { + _commonStats.closes++; + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> LimitSkipStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->specific = std::make_unique<LimitSkipStats>(_specificStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* LimitSkipStage::getSpecificStats() const { + return &_specificStats; +} + +std::vector<DebugPrinter::Block> LimitSkipStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + if (!_skip) { + DebugPrinter::addKeyword(ret, "limit"); + ret.emplace_back(std::to_string(*_limit)); + } else { + DebugPrinter::addKeyword(ret, "limitskip"); + ret.emplace_back(_limit ? std::to_string(*_limit) : "none"); + ret.emplace_back(std::to_string(*_skip)); + } + DebugPrinter::addNewLine(ret); + + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/limit_skip.h b/src/mongo/db/exec/sbe/stages/limit_skip.h new file mode 100644 index 00000000000..9fb285d1648 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/limit_skip.h @@ -0,0 +1,61 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" + +namespace mongo::sbe { +class LimitSkipStage final : public PlanStage { +public: + LimitSkipStage(std::unique_ptr<PlanStage> input, + boost::optional<long long> limit, + boost::optional<long long> skip); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + const boost::optional<long long> _limit; + const boost::optional<long long> _skip; + long long _current; + bool _isEOF; + LimitSkipStats _specificStats; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/loop_join.cpp b/src/mongo/db/exec/sbe/stages/loop_join.cpp new file mode 100644 index 00000000000..85b399cb03e --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/loop_join.cpp @@ -0,0 +1,213 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/loop_join.h" + +#include "mongo/util/str.h" + +namespace mongo::sbe { +LoopJoinStage::LoopJoinStage(std::unique_ptr<PlanStage> outer, + std::unique_ptr<PlanStage> inner, + value::SlotVector outerProjects, + value::SlotVector outerCorrelated, + std::unique_ptr<EExpression> predicate) + : PlanStage("nlj"_sd), + _outerProjects(std::move(outerProjects)), + _outerCorrelated(std::move(outerCorrelated)), + _predicate(std::move(predicate)) { + _children.emplace_back(std::move(outer)); + _children.emplace_back(std::move(inner)); +} + +std::unique_ptr<PlanStage> LoopJoinStage::clone() const { + return std::make_unique<LoopJoinStage>(_children[0]->clone(), + _children[1]->clone(), + _outerProjects, + _outerCorrelated, + _predicate ? _predicate->clone() : nullptr); +} + +void LoopJoinStage::prepare(CompileCtx& ctx) { + for (auto& f : _outerProjects) { + auto [it, inserted] = _outerRefs.emplace(f); + uassert(4822820, str::stream() << "duplicate field: " << f, inserted); + } + _children[0]->prepare(ctx); + + for (auto& f : _outerCorrelated) { + ctx.pushCorrelated(f, _children[0]->getAccessor(ctx, f)); + } + _children[1]->prepare(ctx); + + for (size_t idx = 0; idx < _outerCorrelated.size(); ++idx) { + ctx.popCorrelated(); + } + + if (_predicate) { + ctx.root = this; + _predicateCode = _predicate->compile(ctx); + } +} + +value::SlotAccessor* LoopJoinStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_outerRefs.count(slot)) { + return _children[0]->getAccessor(ctx, slot); + } + + return _children[1]->getAccessor(ctx, slot); +} + +void LoopJoinStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + _outerGetNext = true; + // Do not open the inner child as we do not have values of correlated parameters yet. + // The values are available only after we call getNext on the outer side. +} + +void LoopJoinStage::openInner() { + // (re)open the inner side as it can see the correlated value now. + _children[1]->open(_reOpenInner); + _reOpenInner = true; +} + +PlanState LoopJoinStage::getNext() { + if (_outerGetNext) { + auto state = _children[0]->getNext(); + if (state != PlanState::ADVANCED) { + return trackPlanState(state); + } + + openInner(); + _outerGetNext = false; + } + + for (;;) { + auto state = PlanState::IS_EOF; + bool pass = false; + + do { + + state = _children[1]->getNext(); + if (state == PlanState::ADVANCED) { + if (!_predicateCode) { + pass = true; + } else { + pass = _bytecode.runPredicate(_predicateCode.get()); + } + } + } while (state == PlanState::ADVANCED && !pass); + + if (state == PlanState::ADVANCED) { + return trackPlanState(PlanState::ADVANCED); + } + invariant(state == PlanState::IS_EOF); + + state = _children[0]->getNext(); + if (state != PlanState::ADVANCED) { + return trackPlanState(state); + } + + openInner(); + } +} + +void LoopJoinStage::close() { + _commonStats.closes++; + + if (_reOpenInner) { + _children[1]->close(); + + _reOpenInner = false; + } + + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> LoopJoinStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + ret->children.emplace_back(_children[1]->getStats()); + return ret; +} + +const SpecificStats* LoopJoinStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> LoopJoinStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "nlj"); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _outerProjects.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _outerProjects[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _outerCorrelated.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _outerCorrelated[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + if (_predicate) { + ret.emplace_back("{`"); + DebugPrinter::addBlocks(ret, _predicate->debugPrint()); + ret.emplace_back("`}"); + } + + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + + DebugPrinter::addKeyword(ret, "left"); + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + + DebugPrinter::addKeyword(ret, "right"); + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + DebugPrinter::addBlocks(ret, _children[1]->debugPrint()); + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/loop_join.h b/src/mongo/db/exec/sbe/stages/loop_join.h new file mode 100644 index 00000000000..bf19c50b8f2 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/loop_join.h @@ -0,0 +1,77 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/vm/vm.h" + +namespace mongo::sbe { +class LoopJoinStage final : public PlanStage { +public: + LoopJoinStage(std::unique_ptr<PlanStage> outer, + std::unique_ptr<PlanStage> inner, + value::SlotVector outerProjects, + value::SlotVector outerCorrelated, + std::unique_ptr<EExpression> predicate); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + // Set of variables coming from the outer side. + const value::SlotVector _outerProjects; + // Set of correlated variables from the outer side that are visible on the inner side. They must + // be also present in the _outerProjects. + const value::SlotVector _outerCorrelated; + // If not set then this is a cross product. + const std::unique_ptr<EExpression> _predicate; + + value::SlotSet _outerRefs; + + std::vector<value::SlotAccessor*> _correlatedAccessors; + std::unique_ptr<vm::CodeFragment> _predicateCode; + + vm::ByteCode _bytecode; + bool _reOpenInner{false}; + bool _outerGetNext{false}; + + void openInner(); +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/makeobj.cpp b/src/mongo/db/exec/sbe/stages/makeobj.cpp new file mode 100644 index 00000000000..507458ee152 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/makeobj.cpp @@ -0,0 +1,252 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/makeobj.h" + +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/util/str.h" + +namespace mongo::sbe { +MakeObjStage::MakeObjStage(std::unique_ptr<PlanStage> input, + value::SlotId objSlot, + boost::optional<value::SlotId> rootSlot, + std::vector<std::string> restrictFields, + std::vector<std::string> projectFields, + value::SlotVector projectVars, + bool forceNewObject, + bool returnOldObject) + : PlanStage("mkobj"_sd), + _objSlot(objSlot), + _rootSlot(rootSlot), + _restrictFields(std::move(restrictFields)), + _projectFields(std::move(projectFields)), + _projectVars(std::move(projectVars)), + _forceNewObject(forceNewObject), + _returnOldObject(returnOldObject) { + _children.emplace_back(std::move(input)); + invariant(_projectVars.size() == _projectFields.size()); +} + +std::unique_ptr<PlanStage> MakeObjStage::clone() const { + return std::make_unique<MakeObjStage>(_children[0]->clone(), + _objSlot, + _rootSlot, + _restrictFields, + _projectFields, + _projectVars, + _forceNewObject, + _returnOldObject); +} + +void MakeObjStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + if (_rootSlot) { + _root = _children[0]->getAccessor(ctx, *_rootSlot); + } + for (auto& p : _restrictFields) { + if (p.empty()) { + _restrictAllFields = true; + } else { + auto [it, inserted] = _restrictFieldsSet.emplace(p); + uassert(4822818, str::stream() << "duplicate field: " << p, inserted); + } + } + for (size_t idx = 0; idx < _projectFields.size(); ++idx) { + auto& p = _projectFields[idx]; + + auto [it, inserted] = _projectFieldsMap.emplace(p, idx); + uassert(4822819, str::stream() << "duplicate field: " << p, inserted); + _projects.emplace_back(p, _children[0]->getAccessor(ctx, _projectVars[idx])); + } + _compiled = true; +} + +value::SlotAccessor* MakeObjStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_compiled && slot == _objSlot) { + return &_obj; + } else { + return _children[0]->getAccessor(ctx, slot); + } +} + +void MakeObjStage::projectField(value::Object* obj, size_t idx) { + const auto& p = _projects[idx]; + + auto [tag, val] = p.second->getViewOfValue(); + if (tag != value::TypeTags::Nothing) { + auto [tagCopy, valCopy] = value::copyValue(tag, val); + obj->push_back(p.first, tagCopy, valCopy); + } +} + +void MakeObjStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); +} + +PlanState MakeObjStage::getNext() { + auto state = _children[0]->getNext(); + + if (state == PlanState::ADVANCED) { + auto [tag, val] = value::makeNewObject(); + auto obj = value::getObjectView(val); + absl::flat_hash_set<size_t> alreadyProjected; + + _obj.reset(tag, val); + + if (_root) { + auto [tag, val] = _root->getViewOfValue(); + + if (tag == value::TypeTags::bsonObject) { + auto be = value::bitcastTo<const char*>(val); + auto size = ConstDataView(be).read<LittleEndian<uint32_t>>(); + auto end = be + size; + // Simple heuristic to determine number of fields. + obj->reserve(size / 16); + // Skip document length. + be += 4; + while (*be != 0) { + auto sv = bson::fieldNameView(be); + + if (auto it = _projectFieldsMap.find(sv); + !isFieldRestricted(sv) && it == _projectFieldsMap.end()) { + auto [tag, val] = bson::convertFrom(true, be, end, sv.size()); + auto [copyTag, copyVal] = value::copyValue(tag, val); + obj->push_back(sv, copyTag, copyVal); + } else if (it != _projectFieldsMap.end()) { + projectField(obj, it->second); + alreadyProjected.insert(it->second); + } + + be = bson::advance(be, sv.size()); + } + } else if (tag == value::TypeTags::Object) { + auto objRoot = value::getObjectView(val); + obj->reserve(objRoot->size()); + for (size_t idx = 0; idx < objRoot->size(); ++idx) { + std::string_view sv(objRoot->field(idx)); + + if (auto it = _projectFieldsMap.find(sv); + !isFieldRestricted(sv) && it == _projectFieldsMap.end()) { + auto [tag, val] = objRoot->getAt(idx); + auto [copyTag, copyVal] = value::copyValue(tag, val); + obj->push_back(sv, copyTag, copyVal); + } else if (it != _projectFieldsMap.end()) { + projectField(obj, it->second); + alreadyProjected.insert(it->second); + } + } + } else { + for (size_t idx = 0; idx < _projects.size(); ++idx) { + if (alreadyProjected.count(idx) == 0) { + projectField(obj, idx); + } + } + // If the result is non empty object return it. + if (obj->size() || _forceNewObject) { + return trackPlanState(state); + } + // Now we have to make a decision - return Nothing or the original _root. + if (!_returnOldObject) { + _obj.reset(false, value::TypeTags::Nothing, 0); + } else { + // _root is not an object return it unmodified. + _obj.reset(false, tag, val); + } + return trackPlanState(state); + } + } + for (size_t idx = 0; idx < _projects.size(); ++idx) { + if (alreadyProjected.count(idx) == 0) { + projectField(obj, idx); + } + } + } + return trackPlanState(state); +} + +void MakeObjStage::close() { + _commonStats.closes++; + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> MakeObjStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* MakeObjStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> MakeObjStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "mkobj"); + + DebugPrinter::addIdentifier(ret, _objSlot); + + if (_rootSlot) { + DebugPrinter::addIdentifier(ret, *_rootSlot); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _restrictFields.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _restrictFields[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + } + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _projectFields.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _projectFields[idx]); + ret.emplace_back("="); + DebugPrinter::addIdentifier(ret, _projectVars[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(_forceNewObject ? "true" : "false"); + ret.emplace_back(_returnOldObject ? "true" : "false"); + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/makeobj.h b/src/mongo/db/exec/sbe/stages/makeobj.h new file mode 100644 index 00000000000..94e3c462b92 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/makeobj.h @@ -0,0 +1,86 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/values/value.h" + +namespace mongo::sbe { +class MakeObjStage final : public PlanStage { +public: + MakeObjStage(std::unique_ptr<PlanStage> input, + value::SlotId objSlot, + boost::optional<value::SlotId> rootSlot, + std::vector<std::string> restrictFields, + std::vector<std::string> projectFields, + value::SlotVector projectVars, + bool forceNewObject, + bool returnOldObject); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + void projectField(value::Object* obj, size_t idx); + + bool isFieldRestricted(const std::string_view& sv) const { + return _restrictAllFields || _restrictFieldsSet.count(sv) != 0; + } + + const value::SlotId _objSlot; + const boost::optional<value::SlotId> _rootSlot; + const std::vector<std::string> _restrictFields; + const std::vector<std::string> _projectFields; + const value::SlotVector _projectVars; + const bool _forceNewObject; + const bool _returnOldObject; + + absl::flat_hash_set<std::string> _restrictFieldsSet; + absl::flat_hash_map<std::string, size_t> _projectFieldsMap; + + std::vector<std::pair<std::string, value::SlotAccessor*>> _projects; + + value::OwnedValueAccessor _obj; + + value::SlotAccessor* _root{nullptr}; + + bool _compiled{false}; + bool _restrictAllFields{false}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/plan_stats.cpp b/src/mongo/db/exec/sbe/stages/plan_stats.cpp new file mode 100644 index 00000000000..e17ebe4f0e8 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/plan_stats.cpp @@ -0,0 +1,63 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/plan_stats.h" + +#include <queue> + +namespace mongo::sbe { +size_t calculateNumberOfReads(const PlanStageStats* root) { + size_t numReads{0}; + std::queue<const sbe::PlanStageStats*> remaining; + remaining.push(root); + + while (!remaining.empty()) { + auto stats = remaining.front(); + remaining.pop(); + + if (!stats) { + continue; + } + + if (auto scanStats = dynamic_cast<sbe::ScanStats*>(stats->specific.get())) { + numReads += scanStats->numReads; + } else if (auto indexScanStats = + dynamic_cast<sbe::IndexScanStats*>(stats->specific.get())) { + numReads += indexScanStats->numReads; + } + + for (auto&& child : stats->children) { + remaining.push(child.get()); + } + } + return numReads; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/plan_stats.h b/src/mongo/db/exec/sbe/stages/plan_stats.h new file mode 100644 index 00000000000..d1da4f99699 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/plan_stats.h @@ -0,0 +1,107 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/plan_stats.h" + +namespace mongo::sbe { +struct CommonStats { + CommonStats() = delete; + CommonStats(StringData stageType) : stageType{stageType} {} + + uint64_t estimateObjectSizeInBytes() const { + return sizeof(*this); + } + + StringData stageType; + size_t advances{0}; + size_t opens{0}; + size_t closes{0}; + size_t yields{0}; + size_t unyields{0}; + bool isEOF{false}; +}; +using PlanStageStats = BasePlanStageStats<CommonStats>; + +struct ScanStats : public SpecificStats { + SpecificStats* clone() const final { + return new ScanStats(*this); + } + + uint64_t estimateObjectSizeInBytes() const { + return sizeof(*this); + } + + size_t numReads{0}; +}; + +struct IndexScanStats : public SpecificStats { + SpecificStats* clone() const final { + return new IndexScanStats(*this); + } + + uint64_t estimateObjectSizeInBytes() const { + return sizeof(*this); + } + + size_t numReads{0}; +}; + +struct FilterStats : public SpecificStats { + SpecificStats* clone() const final { + return new FilterStats(*this); + } + + uint64_t estimateObjectSizeInBytes() const { + return sizeof(*this); + } + + size_t numTested{0}; +}; + +struct LimitSkipStats : public SpecificStats { + SpecificStats* clone() const final { + return new LimitSkipStats(*this); + } + + uint64_t estimateObjectSizeInBytes() const { + return sizeof(*this); + } + + boost::optional<long long> limit; + boost::optional<long long> skip; +}; + +/** + * Calculates the total number of physical reads in the given plan stats tree. If a stage can do + * a physical read (e.g. COLLSCAN or IXSCAN), then its 'numReads' stats is added to the total. + */ +size_t calculateNumberOfReads(const PlanStageStats* root); +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/project.cpp b/src/mongo/db/exec/sbe/stages/project.cpp new file mode 100644 index 00000000000..b2e7efe7dc9 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/project.cpp @@ -0,0 +1,129 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/project.h" + +namespace mongo { +namespace sbe { +ProjectStage::ProjectStage(std::unique_ptr<PlanStage> input, + value::SlotMap<std::unique_ptr<EExpression>> projects) + : PlanStage("project"_sd), _projects(std::move(projects)) { + _children.emplace_back(std::move(input)); +} + +std::unique_ptr<PlanStage> ProjectStage::clone() const { + value::SlotMap<std::unique_ptr<EExpression>> projects; + for (auto& [k, v] : _projects) { + projects.emplace(k, v->clone()); + } + return std::make_unique<ProjectStage>(_children[0]->clone(), std::move(projects)); +} + +void ProjectStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + // Compile project expressions here. + for (auto& [slot, expr] : _projects) { + ctx.root = this; + auto code = expr->compile(ctx); + _fields[slot] = {std::move(code), value::OwnedValueAccessor{}}; + } + _compiled = true; +} + +value::SlotAccessor* ProjectStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (auto it = _fields.find(slot); _compiled && it != _fields.end()) { + return &it->second.second; + } else { + return _children[0]->getAccessor(ctx, slot); + } +} +void ProjectStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); +} + +PlanState ProjectStage::getNext() { + auto state = _children[0]->getNext(); + + if (state == PlanState::ADVANCED) { + // Run the project expressions here. + for (auto& p : _fields) { + auto [owned, tag, val] = _bytecode.run(p.second.first.get()); + + // Set the accessors. + p.second.second.reset(owned, tag, val); + } + } + + return trackPlanState(state); +} + +void ProjectStage::close() { + _commonStats.closes++; + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> ProjectStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* ProjectStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> ProjectStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "project"); + + ret.emplace_back("[`"); + bool first = true; + for (auto& p : _projects) { + if (!first) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, p.first); + ret.emplace_back("="); + DebugPrinter::addBlocks(ret, p.second->debugPrint()); + first = false; + } + ret.emplace_back("`]"); + + DebugPrinter::addNewLine(ret); + + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + return ret; +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/project.h b/src/mongo/db/exec/sbe/stages/project.h new file mode 100644 index 00000000000..43011fc27ca --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/project.h @@ -0,0 +1,67 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/vm/vm.h" + +namespace mongo::sbe { +class ProjectStage final : public PlanStage { +public: + ProjectStage(std::unique_ptr<PlanStage> input, + value::SlotMap<std::unique_ptr<EExpression>> projects); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + const value::SlotMap<std::unique_ptr<EExpression>> _projects; + value::SlotMap<std::pair<std::unique_ptr<vm::CodeFragment>, value::OwnedValueAccessor>> _fields; + + vm::ByteCode _bytecode; + + bool _compiled{false}; +}; + +template <typename... Ts> +inline auto makeProjectStage(std::unique_ptr<PlanStage> input, Ts&&... pack) { + return makeS<ProjectStage>(std::move(input), makeEM(std::forward<Ts>(pack)...)); +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp new file mode 100644 index 00000000000..5778e34c3bd --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/scan.cpp @@ -0,0 +1,590 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/scan.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/util/str.h" + +namespace mongo { +namespace sbe { +ScanStage::ScanStage(const NamespaceStringOrUUID& name, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + boost::optional<value::SlotId> seekKeySlot, + bool forward, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker, + ScanOpenCallback openCallback) + : PlanStage(seekKeySlot ? "seek"_sd : "scan"_sd, yieldPolicy), + _name(name), + _recordSlot(recordSlot), + _recordIdSlot(recordIdSlot), + _fields(std::move(fields)), + _vars(std::move(vars)), + _seekKeySlot(seekKeySlot), + _forward(forward), + _tracker(tracker), + _openCallback(openCallback) { + invariant(_fields.size() == _vars.size()); + invariant(!_seekKeySlot || _forward); +} + +std::unique_ptr<PlanStage> ScanStage::clone() const { + return std::make_unique<ScanStage>(_name, + _recordSlot, + _recordIdSlot, + _fields, + _vars, + _seekKeySlot, + _forward, + _yieldPolicy, + _tracker, + _openCallback); +} + +void ScanStage::prepare(CompileCtx& ctx) { + if (_recordSlot) { + _recordAccessor = std::make_unique<value::ViewOfValueAccessor>(); + } + + if (_recordIdSlot) { + _recordIdAccessor = std::make_unique<value::ViewOfValueAccessor>(); + } + + for (size_t idx = 0; idx < _fields.size(); ++idx) { + auto [it, inserted] = + _fieldAccessors.emplace(_fields[idx], std::make_unique<value::ViewOfValueAccessor>()); + uassert(4822814, str::stream() << "duplicate field: " << _fields[idx], inserted); + auto [itRename, insertedRename] = _varAccessors.emplace(_vars[idx], it->second.get()); + uassert(4822815, str::stream() << "duplicate field: " << _vars[idx], insertedRename); + } + + if (_seekKeySlot) { + _seekKeyAccessor = ctx.getAccessor(*_seekKeySlot); + } +} + +value::SlotAccessor* ScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_recordSlot && *_recordSlot == slot) { + return _recordAccessor.get(); + } + + if (_recordIdSlot && *_recordIdSlot == slot) { + return _recordIdAccessor.get(); + } + + if (auto it = _varAccessors.find(slot); it != _varAccessors.end()) { + return it->second; + } + + return ctx.getAccessor(slot); +} + +void ScanStage::doSaveState() { + if (_cursor) { + _cursor->save(); + } + + _coll.reset(); +} + +void ScanStage::doRestoreState() { + invariant(_opCtx); + invariant(!_coll); + + // If this stage is not currently open, then there is nothing to restore. + if (!_open) { + return; + } + + _coll.emplace(_opCtx, _name); + + if (_cursor) { + const bool couldRestore = _cursor->restore(); + uassert(ErrorCodes::CappedPositionLost, + str::stream() + << "CollectionScan died due to position in capped collection being deleted. ", + couldRestore); + } +} + +void ScanStage::doDetachFromOperationContext() { + if (_cursor) { + _cursor->detachFromOperationContext(); + } +} + +void ScanStage::doAttachFromOperationContext(OperationContext* opCtx) { + if (_cursor) { + _cursor->reattachToOperationContext(opCtx); + } +} + +void ScanStage::open(bool reOpen) { + _commonStats.opens++; + invariant(_opCtx); + if (!reOpen) { + invariant(!_cursor); + invariant(!_coll); + _coll.emplace(_opCtx, _name); + } else { + invariant(_cursor); + invariant(_coll); + } + + // TODO: this is currently used only to wait for oplog entries to become visible, so we + // may want to consider to move this logic into storage API instead. + if (_openCallback) { + _openCallback(_opCtx, _coll->getCollection(), reOpen); + } + + if (auto collection = _coll->getCollection()) { + if (_seekKeyAccessor) { + auto [tag, val] = _seekKeyAccessor->getViewOfValue(); + uassert(ErrorCodes::BadValue, + "seek key is wrong type", + tag == value::TypeTags::NumberInt64); + + _key = RecordId{value::bitcastTo<int64_t>(val)}; + } + + if (!_cursor || !_seekKeyAccessor) { + _cursor = collection->getCursor(_opCtx, _forward); + } + } else { + _cursor.reset(); + } + + _open = true; + _firstGetNext = true; +} + +PlanState ScanStage::getNext() { + if (!_cursor) { + return trackPlanState(PlanState::IS_EOF); + } + + checkForInterrupt(_opCtx); + + auto nextRecord = + (_firstGetNext && _seekKeyAccessor) ? _cursor->seekExact(_key) : _cursor->next(); + _firstGetNext = false; + + if (!nextRecord) { + return trackPlanState(PlanState::IS_EOF); + } + + if (_recordAccessor) { + _recordAccessor->reset(value::TypeTags::bsonObject, + value::bitcastFrom<const char*>(nextRecord->data.data())); + } + + if (_recordIdAccessor) { + _recordIdAccessor->reset(value::TypeTags::NumberInt64, + value::bitcastFrom<int64_t>(nextRecord->id.repr())); + } + + if (!_fieldAccessors.empty()) { + auto fieldsToMatch = _fieldAccessors.size(); + auto rawBson = nextRecord->data.data(); + auto be = rawBson + 4; + auto end = rawBson + ConstDataView(rawBson).read<LittleEndian<uint32_t>>(); + for (auto& [name, accessor] : _fieldAccessors) { + accessor->reset(); + } + while (*be != 0) { + auto sv = bson::fieldNameView(be); + if (auto it = _fieldAccessors.find(sv); it != _fieldAccessors.end()) { + // Found the field so convert it to Value. + auto [tag, val] = bson::convertFrom(true, be, end, sv.size()); + + it->second->reset(tag, val); + + if ((--fieldsToMatch) == 0) { + // No need to scan any further so bail out early. + break; + } + } + + be = bson::advance(be, sv.size()); + } + } + + if (_tracker && _tracker->trackProgress<TrialRunProgressTracker::kNumReads>(1)) { + // If we're collecting execution stats during multi-planning and reached the end of the + // trial period (trackProgress() will return 'true' in this case), then we can reset the + // tracker. Note that a trial period is executed only once per a PlanStge tree, and once + // completed never run again on the same tree. + _tracker = nullptr; + } + ++_specificStats.numReads; + return trackPlanState(PlanState::ADVANCED); +} + +void ScanStage::close() { + _commonStats.closes++; + _cursor.reset(); + _coll.reset(); + _open = false; +} + +std::unique_ptr<PlanStageStats> ScanStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->specific = std::make_unique<ScanStats>(_specificStats); + return ret; +} + +const SpecificStats* ScanStage::getSpecificStats() const { + return &_specificStats; +} + +std::vector<DebugPrinter::Block> ScanStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + + if (_seekKeySlot) { + DebugPrinter::addKeyword(ret, "seek"); + + DebugPrinter::addIdentifier(ret, _seekKeySlot.get()); + } else { + DebugPrinter::addKeyword(ret, "scan"); + } + + + if (_recordSlot) { + DebugPrinter::addIdentifier(ret, _recordSlot.get()); + } + + if (_recordIdSlot) { + DebugPrinter::addIdentifier(ret, _recordIdSlot.get()); + } + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _fields.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vars[idx]); + ret.emplace_back("="); + DebugPrinter::addIdentifier(ret, _fields[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back("@\"`"); + DebugPrinter::addIdentifier(ret, _name.toString()); + ret.emplace_back("`\""); + + return ret; +} + +ParallelScanStage::ParallelScanStage(const NamespaceStringOrUUID& name, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + PlanYieldPolicy* yieldPolicy) + : PlanStage("pscan"_sd, yieldPolicy), + _name(name), + _recordSlot(recordSlot), + _recordIdSlot(recordIdSlot), + _fields(std::move(fields)), + _vars(std::move(vars)) { + invariant(_fields.size() == _vars.size()); + + _state = std::make_shared<ParallelState>(); +} + +ParallelScanStage::ParallelScanStage(const std::shared_ptr<ParallelState>& state, + const NamespaceStringOrUUID& name, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + PlanYieldPolicy* yieldPolicy) + : PlanStage("pscan"_sd, yieldPolicy), + _name(name), + _recordSlot(recordSlot), + _recordIdSlot(recordIdSlot), + _fields(std::move(fields)), + _vars(std::move(vars)), + _state(state) { + invariant(_fields.size() == _vars.size()); +} + +std::unique_ptr<PlanStage> ParallelScanStage::clone() const { + return std::make_unique<ParallelScanStage>( + _state, _name, _recordSlot, _recordIdSlot, _fields, _vars, _yieldPolicy); +} + +void ParallelScanStage::prepare(CompileCtx& ctx) { + if (_recordSlot) { + _recordAccessor = std::make_unique<value::ViewOfValueAccessor>(); + } + + if (_recordIdSlot) { + _recordIdAccessor = std::make_unique<value::ViewOfValueAccessor>(); + } + + for (size_t idx = 0; idx < _fields.size(); ++idx) { + auto [it, inserted] = + _fieldAccessors.emplace(_fields[idx], std::make_unique<value::ViewOfValueAccessor>()); + uassert(4822816, str::stream() << "duplicate field: " << _fields[idx], inserted); + auto [itRename, insertedRename] = _varAccessors.emplace(_vars[idx], it->second.get()); + uassert(4822817, str::stream() << "duplicate field: " << _vars[idx], insertedRename); + } +} + +value::SlotAccessor* ParallelScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_recordSlot && *_recordSlot == slot) { + return _recordAccessor.get(); + } + + if (_recordIdSlot && *_recordIdSlot == slot) { + return _recordIdAccessor.get(); + } + + if (auto it = _varAccessors.find(slot); it != _varAccessors.end()) { + return it->second; + } + + return ctx.getAccessor(slot); +} + +void ParallelScanStage::doSaveState() { + if (_cursor) { + _cursor->save(); + } + + _coll.reset(); +} + +void ParallelScanStage::doRestoreState() { + invariant(_opCtx); + invariant(!_coll); + + // If this stage is not currently open, then there is nothing to restore. + if (!_open) { + return; + } + + _coll.emplace(_opCtx, _name); + + if (_cursor) { + const bool couldRestore = _cursor->restore(); + uassert(ErrorCodes::CappedPositionLost, + str::stream() + << "CollectionScan died due to position in capped collection being deleted. ", + couldRestore); + } +} + +void ParallelScanStage::doDetachFromOperationContext() { + if (_cursor) { + _cursor->detachFromOperationContext(); + } +} + +void ParallelScanStage::doAttachFromOperationContext(OperationContext* opCtx) { + if (_cursor) { + _cursor->reattachToOperationContext(opCtx); + } +} + +void ParallelScanStage::open(bool reOpen) { + invariant(_opCtx); + invariant(!reOpen, "parallel scan is not restartable"); + + invariant(!_cursor); + invariant(!_coll); + _coll.emplace(_opCtx, _name); + auto collection = _coll->getCollection(); + + if (collection) { + { + stdx::unique_lock lock(_state->mutex); + if (_state->ranges.empty()) { + auto ranges = collection->getRecordStore()->numRecords(_opCtx) / 10240; + if (ranges < 2) { + _state->ranges.emplace_back(Range{RecordId{}, RecordId{}}); + } else { + if (ranges > 1024) { + ranges = 1024; + } + auto randomCursor = collection->getRecordStore()->getRandomCursor(_opCtx); + invariant(randomCursor); + std::set<RecordId> rids; + while (ranges--) { + auto nextRecord = randomCursor->next(); + if (nextRecord) { + rids.emplace(nextRecord->id); + } + } + RecordId lastid{}; + for (auto id : rids) { + _state->ranges.emplace_back(Range{lastid, id}); + lastid = id; + } + _state->ranges.emplace_back(Range{lastid, RecordId{}}); + } + } + } + + _cursor = collection->getCursor(_opCtx); + } + + _open = true; +} + +boost::optional<Record> ParallelScanStage::nextRange() { + invariant(_cursor); + _currentRange = _state->currentRange.fetchAndAdd(1); + if (_currentRange < _state->ranges.size()) { + _range = _state->ranges[_currentRange]; + + return _range.begin.isNull() ? _cursor->next() : _cursor->seekExact(_range.begin); + } else { + return boost::none; + } +} + +PlanState ParallelScanStage::getNext() { + if (!_cursor) { + _commonStats.isEOF = true; + return PlanState::IS_EOF; + } + + checkForInterrupt(_opCtx); + + boost::optional<Record> nextRecord; + + do { + nextRecord = needsRange() ? nextRange() : _cursor->next(); + if (!nextRecord) { + _commonStats.isEOF = true; + return PlanState::IS_EOF; + } + + if (!_range.end.isNull() && nextRecord->id == _range.end) { + setNeedsRange(); + nextRecord = boost::none; + } + } while (!nextRecord); + + if (_recordAccessor) { + _recordAccessor->reset(value::TypeTags::bsonObject, + value::bitcastFrom<const char*>(nextRecord->data.data())); + } + + if (_recordIdAccessor) { + _recordIdAccessor->reset(value::TypeTags::NumberInt64, + value::bitcastFrom<int64_t>(nextRecord->id.repr())); + } + + + if (!_fieldAccessors.empty()) { + auto fieldsToMatch = _fieldAccessors.size(); + auto rawBson = nextRecord->data.data(); + auto be = rawBson + 4; + auto end = rawBson + ConstDataView(rawBson).read<LittleEndian<uint32_t>>(); + for (auto& [name, accessor] : _fieldAccessors) { + accessor->reset(); + } + while (*be != 0) { + auto sv = bson::fieldNameView(be); + if (auto it = _fieldAccessors.find(sv); it != _fieldAccessors.end()) { + // Found the field so convert it to Value. + auto [tag, val] = bson::convertFrom(true, be, end, sv.size()); + + it->second->reset(tag, val); + + if ((--fieldsToMatch) == 0) { + // No need to scan any further so bail out early. + break; + } + } + + be = bson::advance(be, sv.size()); + } + } + + return PlanState::ADVANCED; +} + +void ParallelScanStage::close() { + _cursor.reset(); + _coll.reset(); + _open = false; +} + +std::unique_ptr<PlanStageStats> ParallelScanStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + return ret; +} + +const SpecificStats* ParallelScanStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> ParallelScanStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "pscan"); + + if (_recordSlot) { + DebugPrinter::addIdentifier(ret, _recordSlot.get()); + } + + if (_recordIdSlot) { + DebugPrinter::addIdentifier(ret, _recordIdSlot.get()); + } + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _fields.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vars[idx]); + ret.emplace_back("="); + DebugPrinter::addIdentifier(ret, _fields[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back("@\"`"); + DebugPrinter::addIdentifier(ret, _name.toString()); + ret.emplace_back("`\""); + + return ret; +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/scan.h b/src/mongo/db/exec/sbe/stages/scan.h new file mode 100644 index 00000000000..5382f4726bb --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/scan.h @@ -0,0 +1,182 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/db_raii.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/exec/trial_run_progress_tracker.h" +#include "mongo/db/storage/record_store.h" + +namespace mongo { +namespace sbe { +using ScanOpenCallback = std::function<void(OperationContext*, const Collection*, bool)>; + +class ScanStage final : public PlanStage { +public: + ScanStage(const NamespaceStringOrUUID& name, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + boost::optional<value::SlotId> seekKeySlot, + bool forward, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker, + ScanOpenCallback openCallback = {}); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +protected: + void doSaveState() override; + void doRestoreState() override; + void doDetachFromOperationContext() override; + void doAttachFromOperationContext(OperationContext* opCtx) override; + +private: + const NamespaceStringOrUUID _name; + const boost::optional<value::SlotId> _recordSlot; + const boost::optional<value::SlotId> _recordIdSlot; + const std::vector<std::string> _fields; + const value::SlotVector _vars; + const boost::optional<value::SlotId> _seekKeySlot; + const bool _forward; + + // If provided, used during a trial run to accumulate certain execution stats. Once the trial + // run is complete, this pointer is reset to nullptr. + TrialRunProgressTracker* _tracker{nullptr}; + + ScanOpenCallback _openCallback; + + std::unique_ptr<value::ViewOfValueAccessor> _recordAccessor; + std::unique_ptr<value::ViewOfValueAccessor> _recordIdAccessor; + + value::FieldAccessorMap _fieldAccessors; + value::SlotAccessorMap _varAccessors; + value::SlotAccessor* _seekKeyAccessor{nullptr}; + + bool _open{false}; + + std::unique_ptr<SeekableRecordCursor> _cursor; + boost::optional<AutoGetCollectionForRead> _coll; + RecordId _key; + bool _firstGetNext{false}; + + ScanStats _specificStats; +}; + +class ParallelScanStage final : public PlanStage { + struct Range { + RecordId begin; + RecordId end; + }; + struct ParallelState { + Mutex mutex = MONGO_MAKE_LATCH("ParallelScanStage::ParallelState::mutex"); + std::vector<Range> ranges; + AtomicWord<size_t> currentRange{0}; + }; + +public: + ParallelScanStage(const NamespaceStringOrUUID& name, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + PlanYieldPolicy* yieldPolicy); + + ParallelScanStage(const std::shared_ptr<ParallelState>& state, + const NamespaceStringOrUUID& name, + boost::optional<value::SlotId> recordSlot, + boost::optional<value::SlotId> recordIdSlot, + std::vector<std::string> fields, + value::SlotVector vars, + PlanYieldPolicy* yieldPolicy); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +protected: + void doSaveState() final; + void doRestoreState() final; + void doDetachFromOperationContext() final; + void doAttachFromOperationContext(OperationContext* opCtx) final; + +private: + boost::optional<Record> nextRange(); + bool needsRange() const { + return _currentRange == std::numeric_limits<std::size_t>::max(); + } + void setNeedsRange() { + _currentRange = std::numeric_limits<std::size_t>::max(); + } + + const NamespaceStringOrUUID _name; + const boost::optional<value::SlotId> _recordSlot; + const boost::optional<value::SlotId> _recordIdSlot; + const std::vector<std::string> _fields; + const value::SlotVector _vars; + + std::shared_ptr<ParallelState> _state; + + std::unique_ptr<value::ViewOfValueAccessor> _recordAccessor; + std::unique_ptr<value::ViewOfValueAccessor> _recordIdAccessor; + + value::FieldAccessorMap _fieldAccessors; + value::SlotAccessorMap _varAccessors; + + size_t _currentRange{std::numeric_limits<std::size_t>::max()}; + Range _range; + + bool _open{false}; + + std::unique_ptr<SeekableRecordCursor> _cursor; + boost::optional<AutoGetCollectionForRead> _coll; +}; +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/sort.cpp b/src/mongo/db/exec/sbe/stages/sort.cpp new file mode 100644 index 00000000000..8557c70c5db --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/sort.cpp @@ -0,0 +1,205 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/sort.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/util/str.h" + +namespace mongo { +namespace sbe { +SortStage::SortStage(std::unique_ptr<PlanStage> input, + value::SlotVector obs, + std::vector<value::SortDirection> dirs, + value::SlotVector vals, + size_t limit, + TrialRunProgressTracker* tracker) + : PlanStage("sort"_sd), + _obs(std::move(obs)), + _dirs(std::move(dirs)), + _vals(std::move(vals)), + _limit(limit), + _st(value::MaterializedRowComparator{_dirs}), + _tracker(tracker) { + _children.emplace_back(std::move(input)); + + invariant(_obs.size() == _dirs.size()); +} + +std::unique_ptr<PlanStage> SortStage::clone() const { + return std::make_unique<SortStage>(_children[0]->clone(), _obs, _dirs, _vals, _limit, _tracker); +} + +void SortStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + value::SlotSet dupCheck; + + size_t counter = 0; + // Process order by fields. + for (auto& slot : _obs) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822812, str::stream() << "duplicate field: " << slot, inserted); + + _inKeyAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + _outAccessors.emplace(slot, std::make_unique<SortKeyAccessor>(_stIt, counter++)); + } + + counter = 0; + // Process value fields. + for (auto& slot : _vals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822813, str::stream() << "duplicate field: " << slot, inserted); + + _inValueAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + _outAccessors.emplace(slot, std::make_unique<SortValueAccessor>(_stIt, counter++)); + } +} + +value::SlotAccessor* SortStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (auto it = _outAccessors.find(slot); it != _outAccessors.end()) { + return it->second.get(); + } + + return ctx.getAccessor(slot); +} + +void SortStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + + value::MaterializedRow keys; + value::MaterializedRow vals; + + while (_children[0]->getNext() == PlanState::ADVANCED) { + keys._fields.reserve(_inKeyAccessors.size()); + vals._fields.reserve(_inValueAccessors.size()); + + for (auto accesor : _inKeyAccessors) { + keys._fields.push_back(value::OwnedValueAccessor{}); + auto [tag, val] = accesor->copyOrMoveValue(); + keys._fields.back().reset(true, tag, val); + } + for (auto accesor : _inValueAccessors) { + vals._fields.push_back(value::OwnedValueAccessor{}); + auto [tag, val] = accesor->copyOrMoveValue(); + vals._fields.back().reset(true, tag, val); + } + + _st.emplace(std::move(keys), std::move(vals)); + if (_st.size() - 1 == _limit) { + _st.erase(--_st.end()); + } + + if (_tracker && _tracker->trackProgress<TrialRunProgressTracker::kNumResults>(1)) { + // If we either hit the maximum number of document to return during the trial run, or + // if we've performed enough physical reads, stop populating the sort heap and bail out + // from the trial run by raising a special exception to signal a runtime planner that + // this candidate plan has completed its trial run early. Note that the sort stage is a + // blocking operation and until all documents are loaded from the child stage and + // sorted, the control is not returned to the runtime planner, so an raising this + // special is mechanism to stop the trial run without affecting the plan stats of the + // higher level stages. + _tracker = nullptr; + _children[0]->close(); + uasserted(ErrorCodes::QueryTrialRunCompleted, "Trial run early exit"); + } + } + + _children[0]->close(); + + _stIt = _st.end(); +} + +PlanState SortStage::getNext() { + if (_stIt == _st.end()) { + _stIt = _st.begin(); + } else { + ++_stIt; + } + + if (_stIt == _st.end()) { + return trackPlanState(PlanState::IS_EOF); + } + + return trackPlanState(PlanState::ADVANCED); +} + +void SortStage::close() { + _commonStats.closes++; + _st.clear(); +} + +std::unique_ptr<PlanStageStats> SortStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* SortStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> SortStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "sort"); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _obs.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _obs[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _vals.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vals[idx]); + } + ret.emplace_back("`]"); + + if (_limit != std::numeric_limits<size_t>::max()) { + ret.emplace_back(std::to_string(_limit)); + } + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/sort.h b/src/mongo/db/exec/sbe/stages/sort.h new file mode 100644 index 00000000000..43d9963485a --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/sort.h @@ -0,0 +1,81 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/trial_run_progress_tracker.h" + +namespace mongo::sbe { +class SortStage final : public PlanStage { +public: + SortStage(std::unique_ptr<PlanStage> input, + value::SlotVector obs, + std::vector<value::SortDirection> dirs, + value::SlotVector vals, + size_t limit, + TrialRunProgressTracker* tracker); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + using TableType = std:: + multimap<value::MaterializedRow, value::MaterializedRow, value::MaterializedRowComparator>; + + using SortKeyAccessor = value::MaterializedRowKeyAccessor<TableType::iterator>; + using SortValueAccessor = value::MaterializedRowValueAccessor<TableType::iterator>; + + const value::SlotVector _obs; + const std::vector<value::SortDirection> _dirs; + const value::SlotVector _vals; + const size_t _limit; + + std::vector<value::SlotAccessor*> _inKeyAccessors; + std::vector<value::SlotAccessor*> _inValueAccessors; + + value::SlotMap<std::unique_ptr<value::SlotAccessor>> _outAccessors; + + TableType _st; + TableType::iterator _stIt; + + // If provided, used during a trial run to accumulate certain execution stats. Once the trial + // run is complete, this pointer is reset to nullptr. + TrialRunProgressTracker* _tracker{nullptr}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/spool.cpp b/src/mongo/db/exec/sbe/stages/spool.cpp new file mode 100644 index 00000000000..8e7928e6d63 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/spool.cpp @@ -0,0 +1,289 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/spool.h" + +namespace mongo::sbe { +SpoolEagerProducerStage::SpoolEagerProducerStage(std::unique_ptr<PlanStage> input, + SpoolId spoolId, + value::SlotVector vals) + : PlanStage{"espool"_sd}, _spoolId{spoolId}, _vals{std::move(vals)} { + _children.emplace_back(std::move(input)); +} + +std::unique_ptr<PlanStage> SpoolEagerProducerStage::clone() const { + return std::make_unique<SpoolEagerProducerStage>(_children[0]->clone(), _spoolId, _vals); +} + +void SpoolEagerProducerStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + if (!_buffer) { + _buffer = ctx.getSpoolBuffer(_spoolId); + } + + value::SlotSet dupCheck; + size_t counter = 0; + + for (auto slot : _vals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822810, str::stream() << "duplicate field: " << slot, inserted); + + _inAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + _outAccessors.emplace( + slot, value::MaterializedRowAccessor<SpoolBuffer>{*_buffer, _bufferIt, counter++}); + } +} + +value::SlotAccessor* SpoolEagerProducerStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (auto it = _outAccessors.find(slot); it != _outAccessors.end()) { + return &it->second; + } + + return ctx.getAccessor(slot); +} + +void SpoolEagerProducerStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + + if (reOpen) { + _buffer->clear(); + } + + value::MaterializedRow vals; + + while (_children[0]->getNext() == PlanState::ADVANCED) { + vals._fields.reserve(_inAccessors.size()); + + for (auto accessor : _inAccessors) { + vals._fields.push_back(value::OwnedValueAccessor{}); + auto [tag, val] = accessor->copyOrMoveValue(); + vals._fields.back().reset(true, tag, val); + } + + _buffer->emplace_back(std::move(vals)); + } + + _children[0]->close(); + _bufferIt = _buffer->size(); +} + +PlanState SpoolEagerProducerStage::getNext() { + if (_bufferIt == _buffer->size()) { + _bufferIt = 0; + } else { + ++_bufferIt; + } + + if (_bufferIt == _buffer->size()) { + return trackPlanState(PlanState::IS_EOF); + } + + return trackPlanState(PlanState::ADVANCED); +} + +void SpoolEagerProducerStage::close() { + _commonStats.closes++; +} + +std::unique_ptr<PlanStageStats> SpoolEagerProducerStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* SpoolEagerProducerStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> SpoolEagerProducerStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "espool"); + + DebugPrinter::addSpoolIdentifier(ret, _spoolId); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _vals.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vals[idx]); + } + ret.emplace_back("`]"); + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + return ret; +} + +SpoolLazyProducerStage::SpoolLazyProducerStage(std::unique_ptr<PlanStage> input, + SpoolId spoolId, + value::SlotVector vals, + std::unique_ptr<EExpression> predicate) + : PlanStage{"lspool"_sd}, + _spoolId{spoolId}, + _vals{std::move(vals)}, + _predicate{std::move(predicate)} { + _children.emplace_back(std::move(input)); +} + +std::unique_ptr<PlanStage> SpoolLazyProducerStage::clone() const { + return std::make_unique<SpoolLazyProducerStage>( + _children[0]->clone(), _spoolId, _vals, _predicate->clone()); +} + +void SpoolLazyProducerStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + if (!_buffer) { + _buffer = ctx.getSpoolBuffer(_spoolId); + } + + if (_predicate) { + ctx.root = this; + _predicateCode = _predicate->compile(ctx); + } + + value::SlotSet dupCheck; + + for (auto slot : _vals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822811, str::stream() << "duplicate field: " << slot, inserted); + + _inAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + _outAccessors.emplace(slot, value::ViewOfValueAccessor{}); + } + + _compiled = true; +} + +value::SlotAccessor* SpoolLazyProducerStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_compiled) { + if (auto it = _outAccessors.find(slot); it != _outAccessors.end()) { + return &it->second; + } + } else { + return _children[0]->getAccessor(ctx, slot); + } + + return ctx.getAccessor(slot); +} + +void SpoolLazyProducerStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + + if (reOpen) { + _buffer->clear(); + } +} + +PlanState SpoolLazyProducerStage::getNext() { + auto state = _children[0]->getNext(); + + if (state == PlanState::ADVANCED) { + auto pass{true}; + + if (_predicateCode) { + pass = _bytecode.runPredicate(_predicateCode.get()); + } + + if (pass) { + // We either haven't got a predicate, or it has passed. In both cases, we need pass + // through the input values, and store them into the buffer. + value::MaterializedRow vals; + vals._fields.reserve(_inAccessors.size()); + + for (size_t idx = 0; idx < _inAccessors.size(); ++idx) { + auto [tag, val] = _inAccessors[idx]->getViewOfValue(); + _outAccessors[_vals[idx]].reset(tag, val); + + vals._fields.push_back(value::OwnedValueAccessor{}); + auto [copyTag, copyVal] = value::copyValue(tag, val); + vals._fields.back().reset(true, copyTag, copyVal); + } + + _buffer->emplace_back(std::move(vals)); + } else { + // Otherwise, just pass through the input values. + for (size_t idx = 0; idx < _inAccessors.size(); ++idx) { + auto [tag, val] = _inAccessors[idx]->getViewOfValue(); + _outAccessors[_vals[idx]].reset(tag, val); + } + } + } + + return trackPlanState(state); +} + +void SpoolLazyProducerStage::close() { + _commonStats.closes++; +} + +std::unique_ptr<PlanStageStats> SpoolLazyProducerStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* SpoolLazyProducerStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> SpoolLazyProducerStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "lspool"); + + DebugPrinter::addSpoolIdentifier(ret, _spoolId); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _vals.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vals[idx]); + } + ret.emplace_back("`]"); + + if (_predicate) { + ret.emplace_back("{`"); + DebugPrinter::addBlocks(ret, _predicate->debugPrint()); + ret.emplace_back("`}"); + } + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/spool.h b/src/mongo/db/exec/sbe/stages/spool.h new file mode 100644 index 00000000000..c064bf5d5c9 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/spool.h @@ -0,0 +1,252 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" + +namespace mongo::sbe { +/** + * This is a Spool PlanStage which retains a copy of all data it reads from its child in a shared + * 'SpoolBuffer', and can later return this data without having to call its child to produce it + * again. + * + * This spool operates in an 'Eager' producer mode. On the call to 'open()' it will read and store + * the entire input from its child into the buffer. On the 'getNext' call it will return data from + * the buffer. + * + * This producer spool can be connected with multiple consumer spools via a shared 'SpoolBuffer'. + * This stage will be responsible for populating the buffer, while consumers will read from the + * buffer once its populated, each using its own read pointer. + */ +class SpoolEagerProducerStage final : public PlanStage { +public: + SpoolEagerProducerStage(std::unique_ptr<PlanStage> input, + SpoolId spoolId, + value::SlotVector vals); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + std::shared_ptr<SpoolBuffer> _buffer{nullptr}; + size_t _bufferIt; + const SpoolId _spoolId; + + const value::SlotVector _vals; + std::vector<value::SlotAccessor*> _inAccessors; + value::SlotMap<value::MaterializedRowAccessor<SpoolBuffer>> _outAccessors; +}; + +/** + * This is a Spool PlanStage which retains a copy of all data it reads from its child in a shared + * 'SpoolBuffer', and can later return this data without having to call its child to produce it + * again. + * + * This spool operates in a 'Lazy' producer mode. In contrast to the 'Eager' producer spool, on the + * call to 'open()' it will _not_ read and populate the buffer. Instead, on the call to 'getNext' + * it will read and store the input into the buffer, and immediately return it to the caller stage. + * + * This producer spool can be connected with multiple consumer spools via a shared 'SpoolBuffer'. + * This stage will be responsible for populating the buffer in a lazy fashion as described above, + * while consumers will read from the buffer (possibly while it's still being populated), each using + * its own read pointer. + * + * This spool can be parameterized with an optional predicate which can be used to filter the input + * and store only portion of input data into the buffer. Filtered out input data is passed through + * without being stored into the buffer. + */ +class SpoolLazyProducerStage final : public PlanStage { +public: + SpoolLazyProducerStage(std::unique_ptr<PlanStage> input, + SpoolId spoolId, + value::SlotVector vals, + std::unique_ptr<EExpression> predicate); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + std::shared_ptr<SpoolBuffer> _buffer{nullptr}; + const SpoolId _spoolId; + + const value::SlotVector _vals; + std::vector<value::SlotAccessor*> _inAccessors; + value::SlotMap<value::ViewOfValueAccessor> _outAccessors; + + std::unique_ptr<EExpression> _predicate; + std::unique_ptr<vm::CodeFragment> _predicateCode; + vm::ByteCode _bytecode; + bool _compiled{false}; +}; + +/** + * This is Spool PlanStage which operates in read-only mode. It doesn't populate its 'SpoolBuffer' + * with the input data (and as such, it doesn't have an input PlanStage) but reads and returns data + * from a shared 'SpoolBuffer' that is populated by another producer spool stage. + * + * This consumer PlanStage can operate as a Stack Spool, in conjunction with a 'Lazy' producer + * spool. In this mode the consumer spool on each call to 'getNext' first deletes the input from + * buffer, remembered on the previous call to 'getNext', and then moves the read pointer to the last + * element in the buffer and returns it. + * + * Since in 'Stack' mode this spool always returns the last input from the buffer, it does not read + * data in the same order as they were added. It will always return the last added input. For + * example, the lazy spool can add values [1,2,3], then the stack consumer spool will read and + * delete 3, then another two values can be added to the buffer [1,2,4,5], then the consumer spool + * will read and delete 5, and so on. + */ +template <bool IsStack> +class SpoolConsumerStage final : public PlanStage { +public: + SpoolConsumerStage(SpoolId spoolId, value::SlotVector vals) + : PlanStage{IsStack ? "sspool"_sd : "cspool"_sd}, + _spoolId{spoolId}, + _vals{std::move(vals)} {} + + std::unique_ptr<PlanStage> clone() const { + return std::make_unique<SpoolConsumerStage<IsStack>>(_spoolId, _vals); + } + + void prepare(CompileCtx& ctx) { + if (!_buffer) { + _buffer = ctx.getSpoolBuffer(_spoolId); + } + + value::SlotSet dupCheck; + size_t counter = 0; + + for (auto slot : _vals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822809, str::stream() << "duplicate field: " << slot, inserted); + + _outAccessors.emplace( + slot, value::MaterializedRowAccessor<SpoolBuffer>{*_buffer, _bufferIt, counter++}); + } + } + + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (auto it = _outAccessors.find(slot); it != _outAccessors.end()) { + return &it->second; + } + + return ctx.getAccessor(slot); + } + + void open(bool reOpen) { + _commonStats.opens++; + _bufferIt = _buffer->size(); + } + + PlanState getNext() { + if constexpr (IsStack) { + if (_bufferIt != _buffer->size()) { + _buffer->erase(_buffer->begin() + _bufferIt); + } + + if (_buffer->size() == 0) { + return trackPlanState(PlanState::IS_EOF); + } + + _bufferIt = _buffer->size() - 1; + } else { + if (_bufferIt == _buffer->size()) { + _bufferIt = 0; + } else { + ++_bufferIt; + } + + if (_bufferIt == _buffer->size()) { + return trackPlanState(PlanState::IS_EOF); + } + } + return trackPlanState(PlanState::ADVANCED); + } + + void close() { + _commonStats.closes++; + } + + std::unique_ptr<PlanStageStats> getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; + } + + const SpecificStats* getSpecificStats() const { + return nullptr; + } + + std::vector<DebugPrinter::Block> debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, IsStack ? "sspool" : "cspool"); + + DebugPrinter::addSpoolIdentifier(ret, _spoolId); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _vals.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + + DebugPrinter::addIdentifier(ret, _vals[idx]); + } + ret.emplace_back("`]"); + + return ret; + } + +private: + std::shared_ptr<SpoolBuffer> _buffer{nullptr}; + size_t _bufferIt; + const SpoolId _spoolId; + + const value::SlotVector _vals; + value::SlotMap<value::MaterializedRowAccessor<SpoolBuffer>> _outAccessors; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/stages.cpp b/src/mongo/db/exec/sbe/stages/stages.cpp new file mode 100644 index 00000000000..423142a2c98 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/stages.cpp @@ -0,0 +1,79 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/stages.h" + +#include "mongo/db/operation_context.h" + +namespace mongo { +namespace sbe { +void CanSwitchOperationContext::detachFromOperationContext() { + invariant(_opCtx); + + for (auto&& child : _stage->_children) { + child->detachFromOperationContext(); + } + + doDetachFromOperationContext(); + _opCtx = nullptr; +} + +void CanSwitchOperationContext::attachFromOperationContext(OperationContext* opCtx) { + invariant(opCtx); + invariant(!_opCtx); + + for (auto&& child : _stage->_children) { + child->attachFromOperationContext(opCtx); + } + + _opCtx = opCtx; + doAttachFromOperationContext(opCtx); +} + +void CanChangeState::saveState() { + _stage->_commonStats.yields++; + for (auto&& child : _stage->_children) { + child->saveState(); + } + + doSaveState(); +} + +void CanChangeState::restoreState() { + _stage->_commonStats.unyields++; + for (auto&& child : _stage->_children) { + child->restoreState(); + } + + doRestoreState(); +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h new file mode 100644 index 00000000000..23f484c142d --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/stages.h @@ -0,0 +1,289 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/plan_stats.h" +#include "mongo/db/exec/sbe/util/debug_print.h" +#include "mongo/db/exec/sbe/values/value.h" +#include "mongo/db/exec/scoped_timer.h" +#include "mongo/db/operation_context.h" +#include "mongo/db/query/plan_yield_policy.h" + +namespace mongo { +namespace sbe { + +struct CompileCtx; +class PlanStage; +enum class PlanState { ADVANCED, IS_EOF }; + +/** + * Provides methods to detach and re-attach to an operation context, which derived classes may + * override to perform additional actions when these events occur. + */ +class CanSwitchOperationContext { +public: + CanSwitchOperationContext(PlanStage* stage) : _stage(stage) { + invariant(_stage); + } + + /** + * Detaches from the OperationContext and releases any storage-engine state. + * + * It is only legal to call this when in a "saved" state. While in the "detached" state, it is + * only legal to call reattachToOperationContext or the destructor. It is not legal to call + * detachFromOperationContext() while already in the detached state. + * + * Propagates to all children, then calls doDetachFromOperationContext(). + */ + void detachFromOperationContext(); + + /** + * Reattaches to the OperationContext and reacquires any storage-engine state. + * + * It is only legal to call this in the "detached" state. On return, the cursor is left in a + * "saved" state, so callers must still call restoreState to use this object. + * + * Propagates to all children, then calls doReattachToOperationContext(). + */ + void attachFromOperationContext(OperationContext* opCtx); + +protected: + // Derived classes can optionally override these methods. + virtual void doDetachFromOperationContext() {} + virtual void doAttachFromOperationContext(OperationContext* opCtx) {} + + OperationContext* _opCtx{nullptr}; + +private: + PlanStage* const _stage; +}; + +/** + * Provides methods to save and restore the state of the object which derives from this class + * when corresponding events are generated as a response to a change in the underlying data source. + * Derived classes may override these methods to perform additional actions when these events occur. + */ +class CanChangeState { +public: + CanChangeState(PlanStage* stage) : _stage(stage) { + invariant(_stage); + } + + /** + * Notifies the stage that the underlying data source may change. + * + * It is illegal to call work() or isEOF() when a stage is in the "saved" state. May be called + * before the first call to open(), before execution of the plan has begun. + * + * Propagates to all children, then calls doSaveState(). + */ + void saveState(); + + /** + * Notifies the stage that underlying data is stable again and prepares for calls to work(). + * + * Can only be called while the stage in is the "saved" state. + * + * Propagates to all children, then calls doRestoreState(). + * + * Throws a UserException on failure to restore due to a conflicting event such as a + * collection drop. May throw a WriteConflictException, in which case the caller may choose to + * retry. + */ + void restoreState(); + +protected: + // Derived classes can optionally override these methods. + virtual void doSaveState() {} + virtual void doRestoreState() {} + +private: + PlanStage* const _stage; +}; + +/** + * Provides methods to obtain execution statistics specific to a plan stage. + */ +class CanTrackStats { +public: + CanTrackStats(StringData stageType) : _commonStats(stageType) {} + + /** + * Returns a tree of stats. If the stage has any children it must propagate the request for + * stats to them. + */ + virtual std::unique_ptr<PlanStageStats> getStats() const = 0; + + /** + * Get stats specific to this stage. Some stages may not have specific stats, in which + * case they return nullptr. The pointer is *not* owned by the caller. + * + * The returned pointer is only valid when the corresponding stage is also valid. + * It must not exist past the stage. If you need the stats to outlive the stage, + * use the getStats(...) method above. + */ + virtual const SpecificStats* getSpecificStats() const = 0; + + /** + * Get the CommonStats for this stage. The pointer is *not* owned by the caller. + * + * The returned pointer is only valid when the corresponding stage is also valid. + * It must not exist past the stage. If you need the stats to outlive the stage, + * use the getStats(...) method above. + */ + const CommonStats* getCommonStats() const { + return &_commonStats; + } + +protected: + PlanState trackPlanState(PlanState state) { + if (state == PlanState::IS_EOF) { + _commonStats.isEOF = true; + } else { + invariant(state == PlanState::ADVANCED); + _commonStats.advances++; + } + return state; + } + + CommonStats _commonStats; +}; + +/** + * Provides a methods which can be used to check if the current operation has been interrupted. + * Maintains an internal state to maintain the interrupt check period. + */ +class CanInterrupt { +public: + /** + * This object will always be responsible for interrupt checking, but it can also optionally be + * responsible for yielding. In order to enable yielding, the caller should pass a non-null + * 'PlanYieldPolicy' pointer. Yielding may be disabled by providing a nullptr. + */ + explicit CanInterrupt(PlanYieldPolicy* yieldPolicy) : _yieldPolicy(yieldPolicy) {} + + /** + * Checks for interrupt if necessary. If yielding has been enabled for this object, then also + * performs a yield if necessary. + */ + void checkForInterrupt(OperationContext* opCtx) { + invariant(opCtx); + + if (--_interruptCounter == 0) { + _interruptCounter = kInterruptCheckPeriod; + opCtx->checkForInterrupt(); + } + + if (_yieldPolicy && _yieldPolicy->shouldYieldOrInterrupt(opCtx)) { + uassertStatusOK(_yieldPolicy->yieldOrInterrupt(opCtx)); + } + } + +protected: + PlanYieldPolicy* const _yieldPolicy{nullptr}; + +private: + static const int kInterruptCheckPeriod = 128; + int _interruptCounter = kInterruptCheckPeriod; +}; + +/** + * This is an abstract base class of all plan stages in SBE. + */ +class PlanStage : public CanSwitchOperationContext, + public CanChangeState, + public CanTrackStats, + public CanInterrupt { +public: + PlanStage(StringData stageType, PlanYieldPolicy* yieldPolicy) + : CanSwitchOperationContext{this}, + CanChangeState{this}, + CanTrackStats{stageType}, + CanInterrupt{yieldPolicy} {} + + explicit PlanStage(StringData stageType) : PlanStage(stageType, nullptr) {} + + virtual ~PlanStage() = default; + + /** + * The idiomatic C++ pattern of object cloning. Plan stages must be fully copyable as every + * thread in parallel execution needs its own private copy. + */ + virtual std::unique_ptr<PlanStage> clone() const = 0; + + /** + * Prepare this SBE PlanStage tree for execution. Must be called once, and must be called + * prior to open(), getNext(), close(), saveState(), or restoreState(), + */ + virtual void prepare(CompileCtx& ctx) = 0; + + /** + * Returns a slot accessor for a given slot id. This method is only called during the prepare + * phase. + */ + virtual value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) = 0; + + /** + * Opens the plan tree and makes it ready for subsequent open(), getNext(), and close() calls. + * The expectation is that a plan stage acquires resources (e.g. memory buffers) during the open + * call and avoids resource acquisition in getNext(). + * + * When reOpen flag is true then the plan stage should reinitizalize already acquired resources + * (e.g. re-hash, re-sort, re-seek, etc). + */ + virtual void open(bool reOpen) = 0; + + /** + * Moves to the next position. If the end is reached then return EOF otherwise ADVANCED. Callers + * are not required to call getNext until EOF. They can stop consuming results at any time. Once + * EOF is reached it will stay at EOF unless reopened. + */ + virtual PlanState getNext() = 0; + + /** + * The mirror method to open(). It releases any acquired resources. + */ + virtual void close() = 0; + + virtual std::vector<DebugPrinter::Block> debugPrint() const = 0; + + friend class CanSwitchOperationContext; + friend class CanChangeState; + +protected: + std::vector<std::unique_ptr<PlanStage>> _children; +}; + +template <typename T, typename... Args> +inline std::unique_ptr<PlanStage> makeS(Args&&... args) { + return std::make_unique<T>(std::forward<Args>(args)...); +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/text_match.cpp b/src/mongo/db/exec/sbe/stages/text_match.cpp new file mode 100644 index 00000000000..765f1c228c1 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/text_match.cpp @@ -0,0 +1,111 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/text_match.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/values/bson.h" + +namespace mongo::sbe { + +std::unique_ptr<PlanStage> TextMatchStage::clone() const { + return makeS<TextMatchStage>( + _children[0]->clone(), _ftsMatcher.query(), _ftsMatcher.spec(), _inputSlot, _outputSlot); +} + +void TextMatchStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + _inValueAccessor = _children[0]->getAccessor(ctx, _inputSlot); +} + +value::SlotAccessor* TextMatchStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (slot == _outputSlot) { + return &_outValueAccessor; + } + + return _children[0]->getAccessor(ctx, slot); +} + +void TextMatchStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); +} + +PlanState TextMatchStage::getNext() { + auto state = _children[0]->getNext(); + + if (state == PlanState::ADVANCED) { + auto&& [typeTag, value] = _inValueAccessor->getViewOfValue(); + uassert(ErrorCodes::Error(4623400), + "textmatch requires input to be an object", + value::isObject(typeTag)); + BSONObj obj; + if (typeTag == value::TypeTags::bsonObject) { + obj = BSONObj{value::bitcastTo<const char*>(value)}; + } else { + BSONObjBuilder builder; + bson::convertToBsonObj(builder, value::getObjectView(value)); + obj = builder.obj(); + } + const auto matchResult = _ftsMatcher.matches(obj); + _outValueAccessor.reset(value::TypeTags::Boolean, matchResult); + } + + return trackPlanState(state); +} + +void TextMatchStage::close() { + _commonStats.closes++; + _children[0]->close(); +} + +std::vector<DebugPrinter::Block> TextMatchStage::debugPrint() const { + // TODO: Add 'textmatch' to the parser so that the debug output can be parsed back to an + // execution plan. + std::vector<DebugPrinter::Block> ret; + + DebugPrinter::addKeyword(ret, "textmatch"); + DebugPrinter::addIdentifier(ret, _inputSlot); + DebugPrinter::addIdentifier(ret, _outputSlot); + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; +} + +std::unique_ptr<PlanStageStats> TextMatchStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/text_match.h b/src/mongo/db/exec/sbe/stages/text_match.h new file mode 100644 index 00000000000..5499ac88ed4 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/text_match.h @@ -0,0 +1,90 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/fts/fts_matcher.h" + +namespace mongo::sbe { + +/** + * Special PlanStage for evaluating an FTSMatcher. Reads a BSON object from 'inputSlot' and passes + * it to the FTSMatcher. Fills out 'outputSlot' with the resulting boolean. If 'inputSlot' contains + * a value of any type other than 'bsonObject', throws a UserException. + * + * TODO: Can this be expressed via string manipulation EExpressions? That would eliminate the need + * for this special stage. + */ +class TextMatchStage final : public PlanStage { +public: + TextMatchStage(std::unique_ptr<PlanStage> inputStage, + const fts::FTSQueryImpl& ftsQuery, + const fts::FTSSpec& ftsSpec, + value::SlotId inputSlot, + value::SlotId outputSlot) + : PlanStage("textmatch"), + _ftsMatcher(ftsQuery, ftsSpec), + _inputSlot(inputSlot), + _outputSlot(outputSlot) { + _children.emplace_back(std::move(inputStage)); + } + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + + void open(bool reOpen) final; + + PlanState getNext() final; + + void close() final; + + std::vector<DebugPrinter::Block> debugPrint() const final; + + std::unique_ptr<PlanStageStats> getStats() const final; + + const SpecificStats* getSpecificStats() const final { + return nullptr; + } + +private: + // Phrase and negated term matcher. + const fts::FTSMatcher _ftsMatcher; + + const value::SlotId _inputSlot; + const value::SlotId _outputSlot; + + value::SlotAccessor* _inValueAccessor; + value::ViewOfValueAccessor _outValueAccessor; +}; + +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/traverse.cpp b/src/mongo/db/exec/sbe/stages/traverse.cpp new file mode 100644 index 00000000000..dbef4d4dcca --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/traverse.cpp @@ -0,0 +1,318 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/traverse.h" + +namespace mongo::sbe { +TraverseStage::TraverseStage(std::unique_ptr<PlanStage> outer, + std::unique_ptr<PlanStage> inner, + value::SlotId inField, + value::SlotId outField, + value::SlotId outFieldInner, + value::SlotVector outerCorrelated, + std::unique_ptr<EExpression> foldExpr, + std::unique_ptr<EExpression> finalExpr, + boost::optional<size_t> nestedArraysDepth) + : PlanStage("traverse"_sd), + _inField(inField), + _outField(outField), + _outFieldInner(outFieldInner), + _correlatedSlots(std::move(outerCorrelated)), + _fold(std::move(foldExpr)), + _final(std::move(finalExpr)), + _nestedArraysDepth(nestedArraysDepth) { + _children.emplace_back(std::move(outer)); + _children.emplace_back(std::move(inner)); + + if (_inField == _outField && (_fold || _final)) { + uasserted(4822808, "in and out field must not match when folding"); + } +} + +std::unique_ptr<PlanStage> TraverseStage::clone() const { + return std::make_unique<TraverseStage>(_children[0]->clone(), + _children[1]->clone(), + _inField, + _outField, + _outFieldInner, + _correlatedSlots, + _fold ? _fold->clone() : nullptr, + _final ? _final->clone() : nullptr); +} + +void TraverseStage::prepare(CompileCtx& ctx) { + // Prepare the outer side as usual. + _children[0]->prepare(ctx); + + // Get the inField (incoming) accessor. + _inFieldAccessor = _children[0]->getAccessor(ctx, _inField); + + // Prepare the accessor for the correlated parameter. + ctx.pushCorrelated(_inField, &_correlatedAccessor); + for (auto slot : _correlatedSlots) { + ctx.pushCorrelated(slot, _children[0]->getAccessor(ctx, slot)); + } + // Prepare the inner side. + _children[1]->prepare(ctx); + + // Get the output from the inner side. + _outFieldInputAccessor = _children[1]->getAccessor(ctx, _outFieldInner); + + if (_fold) { + ctx.root = this; + _foldCode = _fold->compile(ctx); + } + + if (_final) { + ctx.root = this; + _finalCode = _final->compile(ctx); + } + + // Restore correlated parameters. + for (size_t idx = 0; idx < _correlatedSlots.size(); ++idx) { + ctx.popCorrelated(); + } + ctx.popCorrelated(); + + _compiled = true; +} + +value::SlotAccessor* TraverseStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_outField == slot) { + return &_outFieldOutputAccessor; + } + + if (_compiled) { + // After the compilation pass to the 'outer' child. + return _children[0]->getAccessor(ctx, slot); + } else { + // If internal expressions (fold, final) are not compiled yet then they refer to the 'inner' + // child. + return _children[1]->getAccessor(ctx, slot); + } +} + +void TraverseStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + // Do not open the inner child as we do not have values of correlated parameters yet. + // The values are available only after we call getNext on the outer side. +} + +void TraverseStage::openInner(value::TypeTags tag, value::Value val) { + // Set the correlated value. + _correlatedAccessor.reset(tag, val); + + // And (re)open the inner side as it can see the correlated value now. + _children[1]->open(_reOpenInner); + _reOpenInner = true; +} + +PlanState TraverseStage::getNext() { + auto state = _children[0]->getNext(); + if (state != PlanState::ADVANCED) { + return trackPlanState(state); + } + + traverse(_inFieldAccessor, &_outFieldOutputAccessor, 0); + + return trackPlanState(PlanState::ADVANCED); +} + +bool TraverseStage::traverse(value::SlotAccessor* inFieldAccessor, + value::OwnedValueAccessor* outFieldOutputAccessor, + size_t level) { + auto earlyExit = false; + // Get the value. + auto [tag, val] = inFieldAccessor->getViewOfValue(); + + if (value::isArray(tag)) { + // If it is an array then we have to traverse it. + value::ArrayAccessor inArrayAccessor; + inArrayAccessor.reset(tag, val); + value::Array* arrOut{nullptr}; + + if (!_foldCode) { + // Create a fresh new output array. + // TODO if _inField == _outField then we can do implace update of the input array. + auto [tag, val] = value::makeNewArray(); + arrOut = value::getArrayView(val); + outFieldOutputAccessor->reset(true, tag, val); + } else { + outFieldOutputAccessor->reset(false, value::TypeTags::Nothing, 0); + } + + // Loop over all elements of array. + bool firstValue = true; + for (; !inArrayAccessor.atEnd(); inArrayAccessor.advance()) { + auto [tag, val] = inArrayAccessor.getViewOfValue(); + + if (value::isArray(tag)) { + if (_nestedArraysDepth && level + 1 >= *_nestedArraysDepth) { + continue; + } + + // If the current array element is an array itself, traverse it recursively. + value::OwnedValueAccessor outArrayAccessor; + earlyExit = traverse(&inArrayAccessor, &outArrayAccessor, level + 1); + auto [tag, val] = outArrayAccessor.copyOrMoveValue(); + + if (!_foldCode) { + arrOut->push_back(tag, val); + } else { + outFieldOutputAccessor->reset(true, tag, val); + if (earlyExit) { + break; + } + } + } else { + // Otherwise, execute inner side once for every element of the array. + openInner(tag, val); + auto state = _children[1]->getNext(); + + if (state == PlanState::ADVANCED) { + if (!_foldCode) { + // We have to copy (or move optimization) the value to the array + // as by definition all composite values (arrays, objects) own their + // constituents. + auto [tag, val] = _outFieldInputAccessor->copyOrMoveValue(); + arrOut->push_back(tag, val); + } else { + if (firstValue) { + auto [tag, val] = _outFieldInputAccessor->copyOrMoveValue(); + outFieldOutputAccessor->reset(true, tag, val); + firstValue = false; + } else { + // Fold + auto [owned, tag, val] = _bytecode.run(_foldCode.get()); + if (!owned) { + auto [copyTag, copyVal] = value::copyValue(tag, val); + tag = copyTag; + val = copyVal; + } + outFieldOutputAccessor->reset(true, tag, val); + } + } + + // Check early out condition. + if (_finalCode) { + if (_bytecode.runPredicate(_finalCode.get())) { + earlyExit = true; + break; + } + } + } + } + } + } else { + // For non-arrays we simply execute the inner side once. + openInner(tag, val); + auto state = _children[1]->getNext(); + + if (state == PlanState::IS_EOF) { + outFieldOutputAccessor->reset(); + } else { + auto [tag, val] = _outFieldInputAccessor->getViewOfValue(); + // We don't have to copy the value. + outFieldOutputAccessor->reset(false, tag, val); + } + } + + return earlyExit; +} + +void TraverseStage::close() { + _commonStats.closes++; + + if (_reOpenInner) { + _children[1]->close(); + + _reOpenInner = false; + } + + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> TraverseStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + ret->children.emplace_back(_children[1]->getStats()); + return ret; +} + +const SpecificStats* TraverseStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> TraverseStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "traverse"); + + DebugPrinter::addIdentifier(ret, _outField); + DebugPrinter::addIdentifier(ret, _outFieldInner); + DebugPrinter::addIdentifier(ret, _inField); + + if (_correlatedSlots.size()) { + ret.emplace_back("[`"); + for (size_t idx = 0; idx < _correlatedSlots.size(); ++idx) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + DebugPrinter::addIdentifier(ret, _correlatedSlots[idx]); + } + ret.emplace_back("`]"); + } + if (_fold) { + ret.emplace_back("{`"); + DebugPrinter::addBlocks(ret, _fold->debugPrint()); + ret.emplace_back("`}"); + } + + if (_final) { + ret.emplace_back("{`"); + DebugPrinter::addBlocks(ret, _final->debugPrint()); + ret.emplace_back("`}"); + } + + DebugPrinter::addNewLine(ret); + DebugPrinter::addIdentifier(ret, "in"); + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + DebugPrinter::addBlocks(ret, _children[1]->debugPrint()); + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + DebugPrinter::addIdentifier(ret, "from"); + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/traverse.h b/src/mongo/db/exec/sbe/stages/traverse.h new file mode 100644 index 00000000000..fb7555b745e --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/traverse.h @@ -0,0 +1,112 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/vm/vm.h" + +namespace mongo::sbe { +/** + * This is an array traversal operator. If the input value coming from the 'outer' side is an array + * then we execute the 'inner' side exactly once for every element of the array. The results from + * the 'inner' side are then collected into the output array value. The traversal is recursive and + * the structure of nested arrays is preserved (up to optional depth). If the input value is not an + * array then we execute the inner side just once and return the result. + * + * If an optional 'fold' expression is provided then instead of the output array we combine + * individual results into a single output value. Another expression 'final' controls optional + * short-circuiting (a.k.a. early out) logic. + */ +class TraverseStage final : public PlanStage { +public: + TraverseStage(std::unique_ptr<PlanStage> outer, + std::unique_ptr<PlanStage> inner, + value::SlotId inField, + value::SlotId outField, + value::SlotId outFieldInner, + value::SlotVector outerCorrelated, + std::unique_ptr<EExpression> foldExpr, + std::unique_ptr<EExpression> finalExpr, + boost::optional<size_t> nestedArraysDepth = boost::none); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + void openInner(value::TypeTags tag, value::Value val); + bool traverse(value::SlotAccessor* inFieldAccessor, + value::OwnedValueAccessor* outFieldOutputAccessor, + size_t level); + + // The input slot holding value being traversed. + const value::SlotId _inField; + + // The output slot holding result of the traversal. + const value::SlotId _outField; + + // The result of a single iteration of the traversal. + const value::SlotId _outFieldInner; + + // Slots from the 'outer' side that are explicitly accessible on the 'inner' side. + const value::SlotVector _correlatedSlots; + + // Optional folding expression for combining array elements. + const std::unique_ptr<EExpression> _fold; + + // Optional boolean expression controlling short-circuiting of the fold. + const std::unique_ptr<EExpression> _final; + + // Optional nested arrays recursion depth. + const boost::optional<size_t> _nestedArraysDepth; + + value::SlotAccessor* _inFieldAccessor{nullptr}; + value::ViewOfValueAccessor _correlatedAccessor; + value::OwnedValueAccessor _outFieldOutputAccessor; + value::SlotAccessor* _outFieldInputAccessor{nullptr}; + + std::unique_ptr<vm::CodeFragment> _foldCode; + std::unique_ptr<vm::CodeFragment> _finalCode; + + vm::ByteCode _bytecode; + + bool _compiled{false}; + bool _reOpenInner{false}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/union.cpp b/src/mongo/db/exec/sbe/stages/union.cpp new file mode 100644 index 00000000000..17cbcea734f --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/union.cpp @@ -0,0 +1,190 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/union.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" + +namespace mongo::sbe { +UnionStage::UnionStage(std::vector<std::unique_ptr<PlanStage>> inputStages, + std::vector<value::SlotVector> inputVals, + value::SlotVector outputVals) + : PlanStage("union"_sd), _inputVals{std::move(inputVals)}, _outputVals{std::move(outputVals)} { + _children = std::move(inputStages); + + invariant(_children.size() > 0); + invariant(_children.size() == _inputVals.size()); + invariant(std::all_of( + _inputVals.begin(), _inputVals.end(), [size = _outputVals.size()](const auto& slots) { + return slots.size() == size; + })); +} + +std::unique_ptr<PlanStage> UnionStage::clone() const { + std::vector<std::unique_ptr<PlanStage>> inputStages; + for (auto& child : _children) { + inputStages.emplace_back(child->clone()); + } + return std::make_unique<UnionStage>(std::move(inputStages), _inputVals, _outputVals); +} + +void UnionStage::prepare(CompileCtx& ctx) { + value::SlotSet dupCheck; + + for (size_t childNum = 0; childNum < _children.size(); childNum++) { + _children[childNum]->prepare(ctx); + + for (auto slot : _inputVals[childNum]) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822806, str::stream() << "duplicate field: " << slot, inserted); + + _inValueAccessors[_children[childNum].get()].emplace_back( + _children[childNum]->getAccessor(ctx, slot)); + } + } + + for (auto slot : _outputVals) { + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822807, str::stream() << "duplicate field: " << slot, inserted); + + _outValueAccessors.emplace_back(value::ViewOfValueAccessor{}); + } +} + +value::SlotAccessor* UnionStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + for (size_t idx = 0; idx < _outputVals.size(); idx++) { + if (_outputVals[idx] == slot) { + return &_outValueAccessors[idx]; + } + } + + return ctx.getAccessor(slot); +} + +void UnionStage::open(bool reOpen) { + _commonStats.opens++; + if (reOpen) { + std::queue<UnionBranch> emptyQueue; + swap(_remainingBranchesToDrain, emptyQueue); + } + + for (auto& child : _children) { + _remainingBranchesToDrain.push({child.get(), reOpen}); + } + + _remainingBranchesToDrain.front().open(); + _currentStage = _remainingBranchesToDrain.front().stage; +} + +PlanState UnionStage::getNext() { + auto state = PlanState::IS_EOF; + + while (!_remainingBranchesToDrain.empty() && state != PlanState::ADVANCED) { + if (!_currentStage) { + auto& branch = _remainingBranchesToDrain.front(); + branch.open(); + _currentStage = branch.stage; + } + state = _currentStage->getNext(); + + if (state == PlanState::IS_EOF) { + _currentStage = nullptr; + _remainingBranchesToDrain.front().close(); + _remainingBranchesToDrain.pop(); + } else { + const auto& inValueAccessors = _inValueAccessors[_currentStage]; + + for (size_t idx = 0; idx < inValueAccessors.size(); idx++) { + auto [tag, val] = inValueAccessors[idx]->getViewOfValue(); + _outValueAccessors[idx].reset(tag, val); + } + } + } + + return trackPlanState(state); +} + +void UnionStage::close() { + _commonStats.closes++; + _currentStage = nullptr; + while (!_remainingBranchesToDrain.empty()) { + _remainingBranchesToDrain.front().close(); + _remainingBranchesToDrain.pop(); + } +} + +std::unique_ptr<PlanStageStats> UnionStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + for (auto&& child : _children) { + ret->children.emplace_back(child->getStats()); + } + return ret; +} + +const SpecificStats* UnionStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> UnionStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "union"); + + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _outputVals.size(); idx++) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + DebugPrinter::addIdentifier(ret, _outputVals[idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + ret.emplace_back(DebugPrinter::Block::cmdIncIndent); + for (size_t childNum = 0; childNum < _children.size(); childNum++) { + ret.emplace_back(DebugPrinter::Block("[`")); + for (size_t idx = 0; idx < _inputVals[childNum].size(); idx++) { + if (idx) { + ret.emplace_back(DebugPrinter::Block("`,")); + } + DebugPrinter::addIdentifier(ret, _inputVals[childNum][idx]); + } + ret.emplace_back(DebugPrinter::Block("`]")); + + DebugPrinter::addBlocks(ret, _children[childNum]->debugPrint()); + + if (childNum + 1 < _children.size()) { + DebugPrinter::addNewLine(ret); + } + } + ret.emplace_back(DebugPrinter::Block::cmdDecIndent); + + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/union.h b/src/mongo/db/exec/sbe/stages/union.h new file mode 100644 index 00000000000..abd8ce2d9d8 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/union.h @@ -0,0 +1,83 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <queue> + +#include "mongo/db/exec/sbe/stages/stages.h" + +namespace mongo::sbe { +class UnionStage final : public PlanStage { +public: + UnionStage(std::vector<std::unique_ptr<PlanStage>> inputStages, + std::vector<value::SlotVector> inputVals, + value::SlotVector outputVals); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + struct UnionBranch { + PlanStage* stage{nullptr}; + const bool reOpen{false}; + bool isOpen{false}; + + void open() { + if (!isOpen) { + stage->open(reOpen); + isOpen = true; + } + } + + void close() { + if (isOpen) { + stage->close(); + isOpen = false; + } + } + }; + + const std::vector<value::SlotVector> _inputVals; + const value::SlotVector _outputVals; + stdx::unordered_map<PlanStage*, std::vector<value::SlotAccessor*>> _inValueAccessors; + std::vector<value::ViewOfValueAccessor> _outValueAccessors; + std::queue<UnionBranch> _remainingBranchesToDrain; + PlanStage* _currentStage{nullptr}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/unwind.cpp b/src/mongo/db/exec/sbe/stages/unwind.cpp new file mode 100644 index 00000000000..2982e8a6a03 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/unwind.cpp @@ -0,0 +1,175 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/stages/unwind.h" + +#include "mongo/util/str.h" + +namespace mongo::sbe { +UnwindStage::UnwindStage(std::unique_ptr<PlanStage> input, + value::SlotId inField, + value::SlotId outField, + value::SlotId outIndex, + bool preserveNullAndEmptyArrays) + : PlanStage("unwind"_sd), + _inField(inField), + _outField(outField), + _outIndex(outIndex), + _preserveNullAndEmptyArrays(preserveNullAndEmptyArrays) { + _children.emplace_back(std::move(input)); + + if (_outField == _outIndex) { + uasserted(4822805, str::stream() << "duplicate field name: " << _outField); + } +} + +std::unique_ptr<PlanStage> UnwindStage::clone() const { + return std::make_unique<UnwindStage>( + _children[0]->clone(), _inField, _outField, _outIndex, _preserveNullAndEmptyArrays); +} + +void UnwindStage::prepare(CompileCtx& ctx) { + _children[0]->prepare(ctx); + + // Get the inField (incoming) accessor. + _inFieldAccessor = _children[0]->getAccessor(ctx, _inField); + + // Prepare the outField output accessor. + _outFieldOutputAccessor = std::make_unique<value::ViewOfValueAccessor>(); + + // Prepare the outIndex output accessor. + _outIndexOutputAccessor = std::make_unique<value::ViewOfValueAccessor>(); +} + +value::SlotAccessor* UnwindStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { + if (_outField == slot) { + return _outFieldOutputAccessor.get(); + } + + if (_outIndex == slot) { + return _outIndexOutputAccessor.get(); + } + + return _children[0]->getAccessor(ctx, slot); +} + +void UnwindStage::open(bool reOpen) { + _commonStats.opens++; + _children[0]->open(reOpen); + + _index = 0; + _inArray = false; +} + +PlanState UnwindStage::getNext() { + if (!_inArray) { + do { + auto state = _children[0]->getNext(); + if (state != PlanState::ADVANCED) { + return trackPlanState(state); + } + + // Get the value. + auto [tag, val] = _inFieldAccessor->getViewOfValue(); + + if (value::isArray(tag)) { + _inArrayAccessor.reset(tag, val); + _index = 0; + _inArray = true; + + // Empty input array. + if (_inArrayAccessor.atEnd()) { + _inArray = false; + if (_preserveNullAndEmptyArrays) { + _outFieldOutputAccessor->reset(value::TypeTags::Nothing, 0); + _outIndexOutputAccessor->reset(value::TypeTags::NumberInt64, _index); + return trackPlanState(PlanState::ADVANCED); + } + } + } else { + bool nullOrNothing = + tag == value::TypeTags::Null || tag == value::TypeTags::Nothing; + + if (!nullOrNothing || _preserveNullAndEmptyArrays) { + _outFieldOutputAccessor->reset(tag, val); + _outIndexOutputAccessor->reset(value::TypeTags::Nothing, 0); + return trackPlanState(PlanState::ADVANCED); + } + } + } while (!_inArray); + } + + // We are inside the array so pull out the current element and advance. + auto [tagElem, valElem] = _inArrayAccessor.getViewOfValue(); + + _outFieldOutputAccessor->reset(tagElem, valElem); + _outIndexOutputAccessor->reset(value::TypeTags::NumberInt64, _index); + + _inArrayAccessor.advance(); + ++_index; + + if (_inArrayAccessor.atEnd()) { + _inArray = false; + } + + return trackPlanState(PlanState::ADVANCED); +} + +void UnwindStage::close() { + _commonStats.closes++; + _children[0]->close(); +} + +std::unique_ptr<PlanStageStats> UnwindStage::getStats() const { + auto ret = std::make_unique<PlanStageStats>(_commonStats); + ret->children.emplace_back(_children[0]->getStats()); + return ret; +} + +const SpecificStats* UnwindStage::getSpecificStats() const { + return nullptr; +} + +std::vector<DebugPrinter::Block> UnwindStage::debugPrint() const { + std::vector<DebugPrinter::Block> ret; + DebugPrinter::addKeyword(ret, "unwind"); + + DebugPrinter::addIdentifier(ret, _outField); + DebugPrinter::addIdentifier(ret, _outIndex); + DebugPrinter::addIdentifier(ret, _inField); + ret.emplace_back(_preserveNullAndEmptyArrays ? "true" : "false"); + + DebugPrinter::addNewLine(ret); + DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); + + return ret; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/unwind.h b/src/mongo/db/exec/sbe/stages/unwind.h new file mode 100644 index 00000000000..749e22afee9 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/unwind.h @@ -0,0 +1,70 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/stages/stages.h" + +namespace mongo::sbe { +class UnwindStage final : public PlanStage { +public: + UnwindStage(std::unique_ptr<PlanStage> input, + value::SlotId inField, + value::SlotId outField, + value::SlotId outIndex, + bool preserveNullAndEmptyArrays); + + std::unique_ptr<PlanStage> clone() const final; + + void prepare(CompileCtx& ctx) final; + value::SlotAccessor* getAccessor(CompileCtx& ctx, value::SlotId slot) final; + void open(bool reOpen) final; + PlanState getNext() final; + void close() final; + + std::unique_ptr<PlanStageStats> getStats() const final; + const SpecificStats* getSpecificStats() const final; + std::vector<DebugPrinter::Block> debugPrint() const final; + +private: + const value::SlotId _inField; + const value::SlotId _outField; + const value::SlotId _outIndex; + const bool _preserveNullAndEmptyArrays; + + value::SlotAccessor* _inFieldAccessor{nullptr}; + std::unique_ptr<value::ViewOfValueAccessor> _outFieldOutputAccessor; + std::unique_ptr<value::ViewOfValueAccessor> _outIndexOutputAccessor; + + value::ArrayAccessor _inArrayAccessor; + + size_t _index{0}; + bool _inArray{false}; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/util/debug_print.cpp b/src/mongo/db/exec/sbe/util/debug_print.cpp new file mode 100644 index 00000000000..ca7983f38e8 --- /dev/null +++ b/src/mongo/db/exec/sbe/util/debug_print.cpp @@ -0,0 +1,123 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/util/debug_print.h" + +#include "mongo/db/exec/sbe/stages/stages.h" + +namespace mongo { +namespace sbe { +std::string DebugPrinter::print(std::vector<Block> blocks) { + std::string ret; + int ident = 0; + for (auto& b : blocks) { + bool addSpace = true; + switch (b.cmd) { + case Block::cmdIncIndent: + ++ident; + ret.append("\n"); + addIndent(ident, ret); + break; + case Block::cmdDecIndent: + --ident; + ret.append("\n"); + addIndent(ident, ret); + break; + case Block::cmdNewLine: + ret.append("\n"); + addIndent(ident, ret); + break; + case Block::cmdNone: + break; + case Block::cmdNoneNoSpace: + addSpace = false; + break; + case Block::cmdColorRed: + if (_colorConsole) { + ret.append("\033[0;31m"); + } + break; + case Block::cmdColorGreen: + if (_colorConsole) { + ret.append("\033[0;32m"); + } + break; + case Block::cmdColorBlue: + if (_colorConsole) { + ret.append("\033[0;34m"); + } + break; + case Block::cmdColorCyan: + if (_colorConsole) { + ret.append("\033[0;36m"); + } + break; + case Block::cmdColorYellow: + if (_colorConsole) { + ret.append("\033[0;33m"); + } + break; + case Block::cmdColorNone: + if (_colorConsole) { + ret.append("\033[0m"); + } + break; + } + + std::string_view sv(b.str); + if (!sv.empty()) { + if (sv.front() == '`') { + sv.remove_prefix(1); + if (!ret.empty() && ret.back() == ' ') { + ret.resize(ret.size() - 1, 0); + } + } + if (!sv.empty() && sv.back() == '`') { + sv.remove_suffix(1); + addSpace = false; + } + if (!sv.empty()) { + ret.append(sv); + if (addSpace) { + ret.append(" "); + } + } + } + } + + return ret; +} + +std::string DebugPrinter::print(PlanStage* s) { + return print(s->debugPrint()); +} +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/util/debug_print.h b/src/mongo/db/exec/sbe/util/debug_print.h new file mode 100644 index 00000000000..b500905f281 --- /dev/null +++ b/src/mongo/db/exec/sbe/util/debug_print.h @@ -0,0 +1,133 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <string> +#include <vector> + +#include "mongo/db/exec/sbe/values/slot_id_generator.h" +#include "mongo/db/exec/sbe/values/value.h" +#include "mongo/util/str.h" + +namespace mongo { +namespace sbe { +class PlanStage; + +class DebugPrinter { +public: + struct Block { + enum Command { + cmdIncIndent, + cmdDecIndent, + cmdNone, + cmdNoneNoSpace, + cmdNewLine, + cmdColorRed, + cmdColorGreen, + cmdColorBlue, + cmdColorCyan, + cmdColorYellow, + cmdColorNone + }; + + Command cmd; + std::string str; + + Block(std::string_view s) : cmd(cmdNone), str(s) {} + + Block(Command c, std::string_view s) : cmd(c), str(s) {} + + Block(Command c) : cmd(c) {} + }; + + DebugPrinter(bool colorConsole = false) : _colorConsole(colorConsole) {} + + static void addKeyword(std::vector<Block>& ret, std::string_view k) { + ret.emplace_back(Block::cmdColorCyan); + ret.emplace_back(Block{Block::cmdNoneNoSpace, k}); + ret.emplace_back(Block::cmdColorNone); + ret.emplace_back(Block{Block::cmdNoneNoSpace, " "}); + } + + static void addIdentifier(std::vector<Block>& ret, value::SlotId slot) { + std::string name{str::stream() << "s" << slot}; + ret.emplace_back(Block::cmdColorGreen); + ret.emplace_back(Block{Block::cmdNoneNoSpace, name}); + ret.emplace_back(Block::cmdColorNone); + ret.emplace_back(Block{Block::cmdNoneNoSpace, " "}); + } + + static void addIdentifier(std::vector<Block>& ret, FrameId frameId, value::SlotId slot) { + std::string name{str::stream() << "l" << frameId << "." << slot}; + ret.emplace_back(Block::cmdColorGreen); + ret.emplace_back(Block{Block::cmdNoneNoSpace, name}); + ret.emplace_back(Block::cmdColorNone); + ret.emplace_back(Block{Block::cmdNoneNoSpace, " "}); + } + + static void addIdentifier(std::vector<Block>& ret, std::string_view k) { + ret.emplace_back(Block::cmdColorGreen); + ret.emplace_back(Block{Block::cmdNoneNoSpace, k}); + ret.emplace_back(Block::cmdColorNone); + ret.emplace_back(Block{Block::cmdNoneNoSpace, " "}); + } + + static void addSpoolIdentifier(std::vector<Block>& ret, SpoolId spool) { + std::string name{str::stream() << spool}; + ret.emplace_back(Block::cmdColorGreen); + ret.emplace_back(Block{Block::cmdNoneNoSpace, name}); + ret.emplace_back(Block::cmdColorNone); + ret.emplace_back(Block{Block::cmdNoneNoSpace, " "}); + } + + static void addNewLine(std::vector<Block>& ret) { + ret.emplace_back(Block::cmdNewLine); + } + + static void addBlocks(std::vector<Block>& ret, std::vector<Block> blocks) { + ret.insert(ret.end(), + std::make_move_iterator(blocks.begin()), + std::make_move_iterator(blocks.end())); + } + std::string print(PlanStage* s); + +private: + bool _colorConsole; + + void addIndent(int ident, std::string& s) { + for (int i = 0; i < ident; ++i) { + s.append(" "); + } + } + + std::string print(std::vector<Block> blocks); +}; +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/values/bson.cpp b/src/mongo/db/exec/sbe/values/bson.cpp new file mode 100644 index 00000000000..43176a63551 --- /dev/null +++ b/src/mongo/db/exec/sbe/values/bson.cpp @@ -0,0 +1,356 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/values/bson.h" + +namespace mongo { +namespace sbe { +namespace bson { +// clang-format off +static uint8_t advanceTable[] = { + 0xff, // End + 8, // double + 0xff, // string + 0xfe, // document + 0xfe, // document + 0x80, // binary ??? +1 ? + 0, // Undefined(value) - Deprecated + 12, // ObjectId + 1, // Boolean + 8, // UTC datetime + 0, // Null value + 0x80, // Regular expression + 0x80, // DBPointer + 0x80, // JavaScript code + 0x80, // Symbol + 0x80, // JavaScript code w/ scope ???? + 4, // 32-bit integer + 8, // Timestamp + 8, // 64-bit integer + 16 // 128-bit decimal floating point + +}; +// clang-format on + +const char* advance(const char* be, size_t fieldNameSize) { + auto type = static_cast<unsigned char>(*be); + + be += 1 /*type*/ + fieldNameSize + 1 /*zero at the end of fieldname*/; + if (type < sizeof(advanceTable)) { + auto advOffset = advanceTable[type]; + if (advOffset < 128) { + be += advOffset; + } else { + be += ConstDataView(be).read<LittleEndian<uint32_t>>(); + if (advOffset == 0xff) { + be += 4; + } else if (advOffset == 0xfe) { + } else { + if (static_cast<BSONType>(type) == BSONType::BinData) { + be += 5; + } else { + uasserted(4822803, "unsupported bson element"); + } + } + } + } else { + uasserted(4822804, "unsupported bson element"); + } + + return be; +} + +std::pair<value::TypeTags, value::Value> convertFrom(bool view, + const char* be, + const char* end, + size_t fieldNameSize) { + auto type = static_cast<BSONType>(*be); + // Advance the 'be' pointer; + be += 1 + fieldNameSize + 1; + + switch (type) { + case BSONType::NumberDouble: { + auto dbl = ConstDataView(be).read<LittleEndian<double>>(); + return {value::TypeTags::NumberDouble, value::bitcastFrom(dbl)}; + } + case BSONType::NumberDecimal: { + if (view) { + return {value::TypeTags::NumberDecimal, value::bitcastFrom(be)}; + } + + uint64_t low = ConstDataView(be).read<LittleEndian<uint64_t>>(); + uint64_t high = ConstDataView(be + sizeof(uint64_t)).read<LittleEndian<uint64_t>>(); + auto dec = Decimal128{Decimal128::Value({low, high})}; + + return value::makeCopyDecimal(dec); + } + case BSONType::String: { + if (view) { + return {value::TypeTags::bsonString, value::bitcastFrom(be)}; + } + // len includes trailing zero. + auto len = ConstDataView(be).read<LittleEndian<uint32_t>>(); + be += sizeof(len); + if (len < value::kSmallStringThreshold) { + value::Value smallString; + // Copy 8 bytes fast if we have space. + if (be + 8 < end) { + memcpy(&smallString, be, 8); + } else { + memcpy(&smallString, be, len); + } + return {value::TypeTags::StringSmall, smallString}; + } else { + auto str = new char[len]; + memcpy(str, be, len); + return {value::TypeTags::StringBig, value::bitcastFrom(str)}; + } + } + case BSONType::Object: { + if (view) { + return {value::TypeTags::bsonObject, value::bitcastFrom(be)}; + } + // Skip document length. + be += 4; + auto [tag, val] = value::makeNewObject(); + auto obj = value::getObjectView(val); + + while (*be != 0) { + auto sv = bson::fieldNameView(be); + + auto [tag, val] = convertFrom(false, be, end, sv.size()); + obj->push_back(sv, tag, val); + + be = advance(be, sv.size()); + } + return {tag, val}; + } + case BSONType::Array: { + if (view) { + return {value::TypeTags::bsonArray, value::bitcastFrom(be)}; + } + // Skip array length. + be += 4; + auto [tag, val] = value::makeNewArray(); + auto arr = value::getArrayView(val); + + while (*be != 0) { + auto sv = bson::fieldNameView(be); + + auto [tag, val] = convertFrom(false, be, end, sv.size()); + arr->push_back(tag, val); + + be = advance(be, sv.size()); + } + return {tag, val}; + } + case BSONType::jstOID: { + if (view) { + return {value::TypeTags::bsonObjectId, value::bitcastFrom(be)}; + } + auto [tag, val] = value::makeNewObjectId(); + memcpy(value::getObjectIdView(val), be, sizeof(value::ObjectIdType)); + return {tag, val}; + } + case BSONType::Bool: + return {value::TypeTags::Boolean, *(be)}; + case BSONType::Date: { + auto integer = ConstDataView(be).read<LittleEndian<int64_t>>(); + return {value::TypeTags::Date, value::bitcastFrom(integer)}; + } + case BSONType::jstNULL: + return {value::TypeTags::Null, 0}; + case BSONType::NumberInt: { + auto integer = ConstDataView(be).read<LittleEndian<int32_t>>(); + return {value::TypeTags::NumberInt32, value::bitcastFrom(integer)}; + } + case BSONType::bsonTimestamp: { + auto val = ConstDataView(be).read<LittleEndian<uint64_t>>(); + return {value::TypeTags::Timestamp, value::bitcastFrom(val)}; + } + case BSONType::NumberLong: { + auto val = ConstDataView(be).read<LittleEndian<int64_t>>(); + return {value::TypeTags::NumberInt64, value::bitcastFrom(val)}; + } + default: + return {value::TypeTags::Nothing, 0}; + } +} +void convertToBsonObj(BSONArrayBuilder& builder, value::ArrayEnumerator arr) { + for (; !arr.atEnd(); arr.advance()) { + auto [tag, val] = arr.getViewOfValue(); + + switch (tag) { + case value::TypeTags::Nothing: + break; + case value::TypeTags::NumberInt32: + builder.append(value::bitcastTo<int32_t>(val)); + break; + case value::TypeTags::NumberInt64: + builder.append(value::bitcastTo<int64_t>(val)); + break; + case value::TypeTags::NumberDouble: + builder.append(value::bitcastTo<double>(val)); + break; + case value::TypeTags::NumberDecimal: + builder.append(value::bitcastTo<Decimal128>(val)); + break; + case value::TypeTags::Date: + builder.append(Date_t::fromMillisSinceEpoch(value::bitcastTo<int64_t>(val))); + break; + case value::TypeTags::Timestamp: + builder.append(Timestamp(value::bitcastTo<uint64_t>(val))); + break; + case value::TypeTags::Boolean: + builder.append(val != 0); + break; + case value::TypeTags::Null: + builder.appendNull(); + break; + case value::TypeTags::StringSmall: + case value::TypeTags::StringBig: + case value::TypeTags::bsonString: { + auto sv = value::getStringView(tag, val); + builder.append(StringData{sv.data(), sv.size()}); + break; + } + case value::TypeTags::Array: { + BSONArrayBuilder subarrBuilder(builder.subarrayStart()); + convertToBsonObj(subarrBuilder, value::ArrayEnumerator{tag, val}); + subarrBuilder.doneFast(); + break; + } + case value::TypeTags::ArraySet: { + BSONArrayBuilder subarrBuilder(builder.subarrayStart()); + convertToBsonObj(subarrBuilder, value::ArrayEnumerator{tag, val}); + subarrBuilder.doneFast(); + break; + } + case value::TypeTags::Object: { + BSONObjBuilder subobjBuilder(builder.subobjStart()); + convertToBsonObj(subobjBuilder, value::getObjectView(val)); + subobjBuilder.doneFast(); + break; + } + case value::TypeTags::ObjectId: + builder.append(OID::from(value::getObjectIdView(val)->data())); + break; + case value::TypeTags::bsonObject: + builder.append(BSONObj{value::bitcastTo<const char*>(val)}); + break; + case value::TypeTags::bsonArray: + builder.append(BSONArray{BSONObj{value::bitcastTo<const char*>(val)}}); + break; + case value::TypeTags::bsonObjectId: + builder.append(OID::from(value::bitcastTo<const char*>(val))); + break; + default: + MONGO_UNREACHABLE; + } + } +} +void convertToBsonObj(BSONObjBuilder& builder, value::Object* obj) { + for (size_t idx = 0; idx < obj->size(); ++idx) { + auto [tag, val] = obj->getAt(idx); + const auto& name = obj->field(idx); + + switch (tag) { + case value::TypeTags::Nothing: + break; + case value::TypeTags::NumberInt32: + builder.append(name, value::bitcastTo<int32_t>(val)); + break; + case value::TypeTags::NumberInt64: + builder.append(name, value::bitcastTo<int64_t>(val)); + break; + case value::TypeTags::NumberDouble: + builder.append(name, value::bitcastTo<double>(val)); + break; + case value::TypeTags::NumberDecimal: + builder.append(name, value::bitcastTo<Decimal128>(val)); + break; + case value::TypeTags::Date: + builder.append(name, Date_t::fromMillisSinceEpoch(value::bitcastTo<int64_t>(val))); + break; + case value::TypeTags::Timestamp: + builder.append(name, Timestamp(value::bitcastTo<uint64_t>(val))); + break; + case value::TypeTags::Boolean: + builder.append(name, val != 0); + break; + case value::TypeTags::Null: + builder.appendNull(name); + break; + case value::TypeTags::StringSmall: + case value::TypeTags::StringBig: + case value::TypeTags::bsonString: { + auto sv = value::getStringView(tag, val); + builder.append(name, StringData{sv.data(), sv.size()}); + break; + } + case value::TypeTags::Array: { + BSONArrayBuilder subarrBuilder(builder.subarrayStart(name)); + convertToBsonObj(subarrBuilder, value::ArrayEnumerator{tag, val}); + subarrBuilder.doneFast(); + break; + } + case value::TypeTags::ArraySet: { + BSONArrayBuilder subarrBuilder(builder.subarrayStart(name)); + convertToBsonObj(subarrBuilder, value::ArrayEnumerator{tag, val}); + subarrBuilder.doneFast(); + break; + } + case value::TypeTags::Object: { + BSONObjBuilder subobjBuilder(builder.subobjStart(name)); + convertToBsonObj(subobjBuilder, value::getObjectView(val)); + subobjBuilder.doneFast(); + break; + } + case value::TypeTags::ObjectId: + builder.append(name, OID::from(value::getObjectIdView(val)->data())); + break; + case value::TypeTags::bsonObject: + builder.appendObject(name, value::bitcastTo<const char*>(val)); + break; + case value::TypeTags::bsonArray: + builder.appendArray(name, BSONObj{value::bitcastTo<const char*>(val)}); + break; + case value::TypeTags::bsonObjectId: + builder.append(name, OID::from(value::bitcastTo<const char*>(val))); + break; + default: + MONGO_UNREACHABLE; + } + } +} +} // namespace bson +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/values/bson.h b/src/mongo/db/exec/sbe/values/bson.h new file mode 100644 index 00000000000..70f87ec204c --- /dev/null +++ b/src/mongo/db/exec/sbe/values/bson.h @@ -0,0 +1,51 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/bson/bsonobjbuilder.h" +#include "mongo/db/exec/sbe/values/value.h" + +namespace mongo { +namespace sbe { +namespace bson { +std::pair<value::TypeTags, value::Value> convertFrom(bool view, + const char* be, + const char* end, + size_t fieldNameSize); +const char* advance(const char* be, size_t fieldNameSize); + +inline auto fieldNameView(const char* be) noexcept { + return std::string_view{be + 1}; +} + +void convertToBsonObj(BSONObjBuilder& builder, value::Object* obj); +} // namespace bson +} // namespace sbe +} // namespace mongo
\ No newline at end of file diff --git a/src/mongo/db/exec/sbe/values/slot_id_generator.h b/src/mongo/db/exec/sbe/values/slot_id_generator.h new file mode 100644 index 00000000000..7fb6c01abe0 --- /dev/null +++ b/src/mongo/db/exec/sbe/values/slot_id_generator.h @@ -0,0 +1,75 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/exec/sbe/values/value.h" + +namespace mongo::sbe::value { +/** + * A reusable id generator suitable for use with integer ids that generates each new id by adding an + * increment to the previously generated id. This generator is not thread safe; calls to + * generateByIncrementing must be serialized. + */ +template <class T> +class IncrementingIdGenerator { +protected: + /** + * Constructs a new generator using 'startingId' as the first generated id and 'incrementStep' + * as the value to add to generate subsequent ids. Note that 'incrementStep' may be negative but + * must not be zero. + */ + IncrementingIdGenerator(T startingId, T incrementStep) + : _currentId(startingId), _incrementStep(incrementStep) {} + + T generateByIncrementing() { + _currentId += _incrementStep; + return _currentId; + } + +private: + T _currentId; + T _incrementStep; +}; + +template <class T> +class IdGenerator : IncrementingIdGenerator<T> { +public: + IdGenerator(T startingId = 0, T incrementStep = 1) + : IncrementingIdGenerator<T>(startingId, incrementStep) {} + + T generate() { + return this->generateByIncrementing(); + } +}; + +using SlotIdGenerator = IdGenerator<value::SlotId>; +using FrameIdGenerator = IdGenerator<FrameId>; +using SpoolIdGenerator = IdGenerator<SpoolId>; +} // namespace mongo::sbe::value diff --git a/src/mongo/db/exec/sbe/values/value.cpp b/src/mongo/db/exec/sbe/values/value.cpp new file mode 100644 index 00000000000..4ed55ffb75c --- /dev/null +++ b/src/mongo/db/exec/sbe/values/value.cpp @@ -0,0 +1,618 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/values/value.h" + +#include <pcrecpp.h> + +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/storage/key_string.h" + +namespace mongo { +namespace sbe { +namespace value { + +std::pair<TypeTags, Value> makeCopyKeyString(const KeyString::Value& inKey) { + auto k = new KeyString::Value(inKey); + return {TypeTags::ksValue, reinterpret_cast<Value>(k)}; +} + +std::pair<TypeTags, Value> makeCopyPcreRegex(const pcrecpp::RE& regex) { + auto ownedRegexVal = sbe::value::bitcastFrom(new pcrecpp::RE(regex)); + return {TypeTags::pcreRegex, ownedRegexVal}; +} + +void releaseValue(TypeTags tag, Value val) noexcept { + switch (tag) { + case TypeTags::NumberDecimal: + delete getDecimalView(val); + break; + case TypeTags::Array: + delete getArrayView(val); + break; + case TypeTags::ArraySet: + delete getArraySetView(val); + break; + case TypeTags::Object: + delete getObjectView(val); + break; + case TypeTags::StringBig: + delete[] getBigStringView(val); + break; + case TypeTags::ObjectId: + delete getObjectIdView(val); + break; + case TypeTags::bsonObject: + case TypeTags::bsonArray: + case TypeTags::bsonObjectId: + delete[] bitcastTo<uint8_t*>(val); + break; + case TypeTags::ksValue: + delete getKeyStringView(val); + break; + case TypeTags::pcreRegex: + delete getPrceRegexView(val); + break; + default: + break; + } +} + +std::ostream& operator<<(std::ostream& os, const TypeTags tag) { + switch (tag) { + case TypeTags::Nothing: + os << "Nothing"; + break; + case TypeTags::NumberInt32: + os << "NumberInt32"; + break; + case TypeTags::NumberInt64: + os << "NumberInt64"; + break; + case TypeTags::NumberDouble: + os << "NumberDouble"; + break; + case TypeTags::NumberDecimal: + os << "NumberDecimal"; + break; + case TypeTags::Date: + os << "Date"; + break; + case TypeTags::Timestamp: + os << "Timestamp"; + break; + case TypeTags::Boolean: + os << "Boolean"; + break; + case TypeTags::Null: + os << "Null"; + break; + case TypeTags::StringSmall: + os << "StringSmall"; + break; + case TypeTags::StringBig: + os << "StringBig"; + break; + case TypeTags::Array: + os << "Array"; + break; + case TypeTags::ArraySet: + os << "ArraySet"; + break; + case TypeTags::Object: + os << "Object"; + break; + case TypeTags::ObjectId: + os << "ObjectId"; + break; + case TypeTags::bsonObject: + os << "bsonObject"; + break; + case TypeTags::bsonArray: + os << "bsonArray"; + break; + case TypeTags::bsonString: + os << "bsonString"; + break; + case TypeTags::bsonObjectId: + os << "bsonObjectId"; + break; + default: + os << "unknown tag"; + break; + } + return os; +} + +void printValue(std::ostream& os, TypeTags tag, Value val) { + switch (tag) { + case value::TypeTags::NumberInt32: + os << bitcastTo<int32_t>(val); + break; + case value::TypeTags::NumberInt64: + os << bitcastTo<int64_t>(val); + break; + case value::TypeTags::NumberDouble: + os << bitcastTo<double>(val); + break; + case value::TypeTags::NumberDecimal: + os << bitcastTo<Decimal128>(val).toString(); + break; + case value::TypeTags::Boolean: + os << ((val) ? "true" : "false"); + break; + case value::TypeTags::Null: + os << "null"; + break; + case value::TypeTags::StringSmall: + os << '"' << getSmallStringView(val) << '"'; + break; + case value::TypeTags::StringBig: + os << '"' << getBigStringView(val) << '"'; + break; + case value::TypeTags::Array: { + auto arr = getArrayView(val); + os << '['; + for (size_t idx = 0; idx < arr->size(); ++idx) { + if (idx != 0) { + os << ", "; + } + auto [tag, val] = arr->getAt(idx); + printValue(os, tag, val); + } + os << ']'; + break; + } + case value::TypeTags::ArraySet: { + auto arr = getArraySetView(val); + os << '['; + bool first = true; + for (const auto& v : arr->values()) { + if (!first) { + os << ", "; + } + first = false; + printValue(os, v.first, v.second); + } + os << ']'; + break; + } + case value::TypeTags::Object: { + auto obj = getObjectView(val); + os << '{'; + for (size_t idx = 0; idx < obj->size(); ++idx) { + if (idx != 0) { + os << ", "; + } + os << '"' << obj->field(idx) << '"'; + os << " : "; + auto [tag, val] = obj->getAt(idx); + printValue(os, tag, val); + } + os << '}'; + break; + } + case value::TypeTags::ObjectId: { + auto objId = getObjectIdView(val); + os << "ObjectId(\"" << OID::from(objId->data()).toString() << "\")"; + break; + } + case value::TypeTags::Nothing: + os << "---===*** NOTHING ***===---"; + break; + case value::TypeTags::bsonArray: { + const char* be = getRawPointerView(val); + const char* end = be + ConstDataView(be).read<LittleEndian<uint32_t>>(); + bool first = true; + // Skip document length. + be += 4; + os << '['; + while (*be != 0) { + auto sv = bson::fieldNameView(be); + + if (!first) { + os << ", "; + } + first = false; + + auto [tag, val] = bson::convertFrom(true, be, end, sv.size()); + printValue(os, tag, val); + + be = bson::advance(be, sv.size()); + } + os << ']'; + break; + } + case value::TypeTags::bsonObject: { + const char* be = getRawPointerView(val); + const char* end = be + ConstDataView(be).read<LittleEndian<uint32_t>>(); + bool first = true; + // Skip document length. + be += 4; + os << '{'; + while (*be != 0) { + auto sv = bson::fieldNameView(be); + + if (!first) { + os << ", "; + } + first = false; + + os << '"' << sv << '"'; + os << " : "; + auto [tag, val] = bson::convertFrom(true, be, end, sv.size()); + printValue(os, tag, val); + + be = bson::advance(be, sv.size()); + } + os << '}'; + break; + } + case value::TypeTags::bsonString: + os << '"' << getStringView(value::TypeTags::bsonString, val) << '"'; + break; + case value::TypeTags::bsonObjectId: + os << "---===*** bsonObjectId ***===---"; + break; + case value::TypeTags::ksValue: { + auto ks = getKeyStringView(val); + os << "KS(" << ks->toString() << ")"; + break; + } + case value::TypeTags::Timestamp: { + Timestamp ts{bitcastTo<uint64_t>(val)}; + os << ts.toString(); + break; + } + case value::TypeTags::pcreRegex: { + auto regex = getPrceRegexView(val); + // TODO: Also include the regex flags. + os << "/" << regex->pattern() << "/"; + break; + } + default: + MONGO_UNREACHABLE; + } +} + +BSONType tagToType(TypeTags tag) noexcept { + switch (tag) { + case TypeTags::Nothing: + return BSONType::EOO; + case TypeTags::NumberInt32: + return BSONType::NumberInt; + case TypeTags::NumberInt64: + return BSONType::NumberLong; + case TypeTags::NumberDouble: + return BSONType::NumberDouble; + case TypeTags::NumberDecimal: + return BSONType::NumberDecimal; + case TypeTags::Date: + return BSONType::Date; + case TypeTags::Timestamp: + return BSONType::bsonTimestamp; + case TypeTags::Boolean: + return BSONType::Bool; + case TypeTags::Null: + return BSONType::jstNULL; + case TypeTags::StringSmall: + return BSONType::String; + case TypeTags::StringBig: + return BSONType::String; + case TypeTags::Array: + return BSONType::Array; + case TypeTags::ArraySet: + return BSONType::Array; + case TypeTags::Object: + return BSONType::Object; + case TypeTags::ObjectId: + return BSONType::jstOID; + case TypeTags::bsonObject: + return BSONType::Object; + case TypeTags::bsonArray: + return BSONType::Array; + case TypeTags::bsonString: + return BSONType::String; + case TypeTags::bsonObjectId: + return BSONType::jstOID; + case TypeTags::ksValue: + // This is completely arbitrary. + return BSONType::EOO; + default: + MONGO_UNREACHABLE; + } +} + +std::size_t hashValue(TypeTags tag, Value val) noexcept { + switch (tag) { + case TypeTags::NumberInt32: + return absl::Hash<int32_t>{}(bitcastTo<int32_t>(val)); + case TypeTags::NumberInt64: + return absl::Hash<int64_t>{}(bitcastTo<int64_t>(val)); + case TypeTags::NumberDouble: + // Force doubles to integers for hashing. + return absl::Hash<int64_t>{}(bitcastTo<double>(val)); + case TypeTags::NumberDecimal: + // Force decimals to integers for hashing. + return absl::Hash<int64_t>{}(bitcastTo<Decimal128>(val).toLong()); + case TypeTags::Date: + return absl::Hash<int64_t>{}(bitcastTo<int64_t>(val)); + case TypeTags::Timestamp: + return absl::Hash<uint64_t>{}(bitcastTo<uint64_t>(val)); + case TypeTags::Boolean: + return val != 0; + case TypeTags::Null: + return 0; + case TypeTags::StringSmall: + case TypeTags::StringBig: + case TypeTags::bsonString: { + auto sv = getStringView(tag, val); + return absl::Hash<std::string_view>{}(sv); + } + case TypeTags::ObjectId: { + auto id = getObjectIdView(val); + return absl::Hash<uint64_t>{}(readFromMemory<uint64_t>(id->data())) ^ + absl::Hash<uint32_t>{}(readFromMemory<uint32_t>(id->data() + 8)); + } + case TypeTags::ksValue: { + return getKeyStringView(val)->hash(); + } + default: + break; + } + + return 0; +} + + +/** + * Performs a three-way comparison for any type that has < and == operators. Additionally, + * guarantees that the result will be exactlty -1, 0, or 1, which is important, because not all + * comparison functions make that guarantee. + * + * The std::string_view::compare(basic_string_view s) function, for example, only promises that it + * will return a value less than 0 in the case that 'this' is less than 's,' whereas we want to + * return exactly -1. + */ +template <typename T> +int32_t compareHelper(const T lhs, const T rhs) noexcept { + return lhs < rhs ? -1 : (lhs == rhs ? 0 : 1); +} + +/* + * Three ways value comparison (aka spacehip operator). + */ +std::pair<TypeTags, Value> compareValue(TypeTags lhsTag, + Value lhsValue, + TypeTags rhsTag, + Value rhsValue) { + if (isNumber(lhsTag) && isNumber(rhsTag)) { + switch (getWidestNumericalType(lhsTag, rhsTag)) { + case TypeTags::NumberInt32: { + auto result = compareHelper(numericCast<int32_t>(lhsTag, lhsValue), + numericCast<int32_t>(rhsTag, rhsValue)); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } + case TypeTags::NumberInt64: { + auto result = compareHelper(numericCast<int64_t>(lhsTag, lhsValue), + numericCast<int64_t>(rhsTag, rhsValue)); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } + case TypeTags::NumberDouble: { + auto result = compareHelper(numericCast<double>(lhsTag, lhsValue), + numericCast<double>(rhsTag, rhsValue)); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } + case TypeTags::NumberDecimal: { + auto result = compareHelper(numericCast<Decimal128>(lhsTag, lhsValue), + numericCast<Decimal128>(rhsTag, rhsValue)); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } + default: + MONGO_UNREACHABLE; + } + } else if (isString(lhsTag) && isString(rhsTag)) { + auto lhsStr = getStringView(lhsTag, lhsValue); + auto rhsStr = getStringView(rhsTag, rhsValue); + auto result = lhsStr.compare(rhsStr); + return {TypeTags::NumberInt32, bitcastFrom(compareHelper(result, 0))}; + } else if (lhsTag == TypeTags::Date && rhsTag == TypeTags::Date) { + auto result = compareHelper(bitcastTo<int64_t>(lhsValue), bitcastTo<int64_t>(rhsValue)); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } else if (lhsTag == TypeTags::Timestamp && rhsTag == TypeTags::Timestamp) { + auto result = compareHelper(bitcastTo<uint64_t>(lhsValue), bitcastTo<uint64_t>(rhsValue)); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } else if (lhsTag == TypeTags::Boolean && rhsTag == TypeTags::Boolean) { + auto result = compareHelper(lhsValue != 0, rhsValue != 0); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } else if (lhsTag == TypeTags::Null && rhsTag == TypeTags::Null) { + return {TypeTags::NumberInt32, 0}; + } else if (isArray(lhsTag) && isArray(rhsTag)) { + auto lhsArr = ArrayEnumerator{lhsTag, lhsValue}; + auto rhsArr = ArrayEnumerator{rhsTag, rhsValue}; + while (!lhsArr.atEnd() && !rhsArr.atEnd()) { + auto [lhsTag, lhsVal] = lhsArr.getViewOfValue(); + auto [rhsTag, rhsVal] = rhsArr.getViewOfValue(); + + auto [tag, val] = compareValue(lhsTag, lhsVal, rhsTag, rhsVal); + if (tag != TypeTags::NumberInt32 || val != 0) { + return {tag, val}; + } + lhsArr.advance(); + rhsArr.advance(); + } + if (lhsArr.atEnd() && rhsArr.atEnd()) { + return {TypeTags::NumberInt32, 0}; + } else if (lhsArr.atEnd()) { + return {TypeTags::NumberInt32, bitcastFrom<int32_t>(-1)}; + } else { + return {TypeTags::NumberInt32, bitcastFrom<int32_t>(1)}; + } + } else if (isObject(lhsTag) && isObject(rhsTag)) { + auto lhsObj = ObjectEnumerator{lhsTag, lhsValue}; + auto rhsObj = ObjectEnumerator{rhsTag, rhsValue}; + while (!lhsObj.atEnd() && !rhsObj.atEnd()) { + auto fieldCmp = lhsObj.getFieldName().compare(rhsObj.getFieldName()); + if (fieldCmp != 0) { + return {TypeTags::NumberInt32, bitcastFrom(compareHelper(fieldCmp, 0))}; + } + + auto [lhsTag, lhsVal] = lhsObj.getViewOfValue(); + auto [rhsTag, rhsVal] = rhsObj.getViewOfValue(); + + auto [tag, val] = compareValue(lhsTag, lhsVal, rhsTag, rhsVal); + if (tag != TypeTags::NumberInt32 || val != 0) { + return {tag, val}; + } + lhsObj.advance(); + rhsObj.advance(); + } + if (lhsObj.atEnd() && rhsObj.atEnd()) { + return {TypeTags::NumberInt32, 0}; + } else if (lhsObj.atEnd()) { + return {TypeTags::NumberInt32, bitcastFrom<int32_t>(-1)}; + } else { + return {TypeTags::NumberInt32, bitcastFrom<int32_t>(1)}; + } + } else if (isObjectId(lhsTag) && isObjectId(rhsTag)) { + auto lhsObjId = lhsTag == TypeTags::ObjectId ? getObjectIdView(lhsValue)->data() + : bitcastTo<uint8_t*>(lhsValue); + auto rhsObjId = rhsTag == TypeTags::ObjectId ? getObjectIdView(rhsValue)->data() + : bitcastTo<uint8_t*>(rhsValue); + auto result = memcmp(lhsObjId, rhsObjId, sizeof(ObjectIdType)); + return {TypeTags::NumberInt32, bitcastFrom(compareHelper(result, 0))}; + } else if (lhsTag == TypeTags::ksValue && rhsTag == TypeTags::ksValue) { + auto result = getKeyStringView(lhsValue)->compare(*getKeyStringView(lhsValue)); + return {TypeTags::NumberInt32, bitcastFrom(result)}; + } else if (lhsTag == TypeTags::Nothing && rhsTag == TypeTags::Nothing) { + // Special case for Nothing in a hash table (group) and sort comparison. + return {TypeTags::NumberInt32, 0}; + } else { + // Different types. + auto result = + canonicalizeBSONType(tagToType(lhsTag)) - canonicalizeBSONType(tagToType(rhsTag)); + invariant(result != 0); + return {TypeTags::NumberInt32, bitcastFrom(compareHelper(result, 0))}; + } +} + +void ArraySet::push_back(TypeTags tag, Value val) { + if (tag != TypeTags::Nothing) { + ValueGuard guard{tag, val}; + auto [it, inserted] = _values.insert({tag, val}); + + if (inserted) { + guard.reset(); + } + } +} + + +std::pair<TypeTags, Value> ArrayEnumerator::getViewOfValue() const { + if (_array) { + return _array->getAt(_index); + } else if (_arraySet) { + return {_iter->first, _iter->second}; + } else { + auto sv = bson::fieldNameView(_arrayCurrent); + return bson::convertFrom(true, _arrayCurrent, _arrayEnd, sv.size()); + } +} + +bool ArrayEnumerator::advance() { + if (_array) { + if (_index < _array->size()) { + ++_index; + } + + return _index < _array->size(); + } else if (_arraySet) { + if (_iter != _arraySet->values().end()) { + ++_iter; + } + + return _iter != _arraySet->values().end(); + } else { + if (*_arrayCurrent != 0) { + auto sv = bson::fieldNameView(_arrayCurrent); + _arrayCurrent = bson::advance(_arrayCurrent, sv.size()); + } + + return *_arrayCurrent != 0; + } +} + +std::pair<TypeTags, Value> ObjectEnumerator::getViewOfValue() const { + if (_object) { + return _object->getAt(_index); + } else { + auto sv = bson::fieldNameView(_objectCurrent); + return bson::convertFrom(true, _objectCurrent, _objectEnd, sv.size()); + } +} + +bool ObjectEnumerator::advance() { + if (_object) { + if (_index < _object->size()) { + ++_index; + } + + return _index < _object->size(); + } else { + if (*_objectCurrent != 0) { + auto sv = bson::fieldNameView(_objectCurrent); + _objectCurrent = bson::advance(_objectCurrent, sv.size()); + } + + return *_objectCurrent != 0; + } +} + +std::string_view ObjectEnumerator::getFieldName() const { + using namespace std::literals; + if (_object) { + if (_index < _object->size()) { + return _object->field(_index); + } else { + return ""sv; + } + } else { + if (*_objectCurrent != 0) { + return bson::fieldNameView(_objectCurrent); + } else { + return ""sv; + } + } +} + +} // namespace value +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/values/value.h b/src/mongo/db/exec/sbe/values/value.h new file mode 100644 index 00000000000..e39e5a7c9ee --- /dev/null +++ b/src/mongo/db/exec/sbe/values/value.h @@ -0,0 +1,1066 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <absl/container/flat_hash_map.h> +#include <absl/container/flat_hash_set.h> +#include <array> +#include <cstdint> +#include <ostream> +#include <string> +#include <utility> +#include <vector> + +#include "mongo/base/data_type_endian.h" +#include "mongo/base/data_view.h" +#include "mongo/platform/decimal128.h" +#include "mongo/util/assert_util.h" + +namespace pcrecpp { +class RE; +} // namespace pcrecpp + +namespace mongo { +/** + * Forward declaration. + */ +namespace KeyString { +class Value; +} +namespace sbe { +using FrameId = int64_t; +using SpoolId = int64_t; + +namespace value { +using SlotId = int64_t; + +/** + * Type dispatch tags. + */ +enum class TypeTags : uint8_t { + // The value does not exist, aka Nothing in the Maybe monad. + Nothing = 0, + + // Numberical data types. + NumberInt32, + NumberInt64, + NumberDouble, + NumberDecimal, + + // Date data types. + Date, + Timestamp, + + Boolean, + Null, + StringSmall, + StringBig, + Array, + ArraySet, + Object, + + ObjectId, + + // TODO add the rest of mongo types (regex, etc.) + + // Raw bson values. + bsonObject, + bsonArray, + bsonString, + bsonObjectId, + + // KeyString::Value + ksValue, + + // Pointer to a compiled PCRE regular expression object. + pcreRegex, +}; + +std::ostream& operator<<(std::ostream& os, const TypeTags tag); + +inline constexpr bool isNumber(TypeTags tag) noexcept { + return tag == TypeTags::NumberInt32 || tag == TypeTags::NumberInt64 || + tag == TypeTags::NumberDouble || tag == TypeTags::NumberDecimal; +} + +inline constexpr bool isString(TypeTags tag) noexcept { + return tag == TypeTags::StringSmall || tag == TypeTags::StringBig || + tag == TypeTags::bsonString; +} + +inline constexpr bool isObject(TypeTags tag) noexcept { + return tag == TypeTags::Object || tag == TypeTags::bsonObject; +} + +inline constexpr bool isArray(TypeTags tag) noexcept { + return tag == TypeTags::Array || tag == TypeTags::ArraySet || tag == TypeTags::bsonArray; +} + +inline constexpr bool isObjectId(TypeTags tag) noexcept { + return tag == TypeTags::ObjectId || tag == TypeTags::bsonObjectId; +} + +/** + * The runtime value. It is a simple 64 bit integer. + */ +using Value = uint64_t; + +/** + * Sort direction of ordered sequence. + */ +enum class SortDirection : uint8_t { Descending, Ascending }; + +/** + * Forward declarations. + */ +void releaseValue(TypeTags tag, Value val) noexcept; +std::pair<TypeTags, Value> copyValue(TypeTags tag, Value val); +void printValue(std::ostream& os, TypeTags tag, Value val); +std::size_t hashValue(TypeTags tag, Value val) noexcept; + +/** + * Three ways value comparison (aka spaceship operator). + */ +std::pair<TypeTags, Value> compareValue(TypeTags lhsTag, + Value lhsValue, + TypeTags rhsTag, + Value rhsValue); + +/** + * RAII guard. + */ +class ValueGuard { +public: + ValueGuard(TypeTags tag, Value val) : _tag(tag), _value(val) {} + ValueGuard() = delete; + ValueGuard(const ValueGuard&) = delete; + ValueGuard(ValueGuard&&) = delete; + ~ValueGuard() { + releaseValue(_tag, _value); + } + + ValueGuard& operator=(const ValueGuard&) = delete; + ValueGuard& operator=(ValueGuard&&) = delete; + + void reset() { + _tag = TypeTags::Nothing; + _value = 0; + } + +private: + TypeTags _tag; + Value _value; +}; + +/** + * This is the SBE representation of objects/documents. It is a relatively simple structure of + * vectors of field names, type tags, and values. + */ +class Object { +public: + Object() = default; + Object(const Object& other) { + // Reserve space in all vectors, they are the same size. We arbitrarily picked _typeTags + // to determine the size. + reserve(other._typeTags.size()); + _names = other._names; + for (size_t idx = 0; idx < other._values.size(); ++idx) { + const auto [tag, val] = copyValue(other._typeTags[idx], other._values[idx]); + _values.push_back(val); + _typeTags.push_back(tag); + } + } + Object(Object&&) = default; + ~Object() { + for (size_t idx = 0; idx < _typeTags.size(); ++idx) { + releaseValue(_typeTags[idx], _values[idx]); + } + } + + void push_back(std::string_view name, TypeTags tag, Value val) { + if (tag != TypeTags::Nothing) { + ValueGuard guard{tag, val}; + // Reserve space in all vectors, they are the same size. We arbitrarily picked _typeTags + // to determine the size. + reserve(_typeTags.size() + 1); + _names.emplace_back(std::string(name)); + + _typeTags.push_back(tag); + _values.push_back(val); + + guard.reset(); + } + } + + std::pair<TypeTags, Value> getField(std::string_view field) { + for (size_t idx = 0; idx < _typeTags.size(); ++idx) { + if (_names[idx] == field) { + return {_typeTags[idx], _values[idx]}; + } + } + return {TypeTags::Nothing, 0}; + } + + auto size() const noexcept { + return _values.size(); + } + + auto& field(size_t idx) const { + return _names[idx]; + } + + std::pair<TypeTags, Value> getAt(std::size_t idx) const { + if (idx >= _values.size()) { + return {TypeTags::Nothing, 0}; + } + + return {_typeTags[idx], _values[idx]}; + } + + void reserve(size_t s) { + // Normalize to at least 1. + s = s ? s : 1; + _typeTags.reserve(s); + _values.reserve(s); + _names.reserve(s); + } + +private: + std::vector<TypeTags> _typeTags; + std::vector<Value> _values; + std::vector<std::string> _names; +}; + +/** + * This is the SBE representation of arrays. It is similar to Object without the field names. + */ +class Array { +public: + Array() = default; + Array(const Array& other) { + // Reserve space in all vectors, they are the same size. We arbitrarily picked _typeTags + // to determine the size. + reserve(other._typeTags.size()); + for (size_t idx = 0; idx < other._values.size(); ++idx) { + const auto [tag, val] = copyValue(other._typeTags[idx], other._values[idx]); + _values.push_back(val); + _typeTags.push_back(tag); + } + } + Array(Array&&) = default; + ~Array() { + for (size_t idx = 0; idx < _typeTags.size(); ++idx) { + releaseValue(_typeTags[idx], _values[idx]); + } + } + + void push_back(TypeTags tag, Value val) { + if (tag != TypeTags::Nothing) { + ValueGuard guard{tag, val}; + // Reserve space in all vectors, they are the same size. We arbitrarily picked _typeTags + // to determine the size. + reserve(_typeTags.size() + 1); + + _typeTags.push_back(tag); + _values.push_back(val); + + guard.reset(); + } + } + + auto size() const noexcept { + return _values.size(); + } + + std::pair<TypeTags, Value> getAt(std::size_t idx) const { + if (idx >= _values.size()) { + return {TypeTags::Nothing, 0}; + } + + return {_typeTags[idx], _values[idx]}; + } + + void reserve(size_t s) { + // Normalize to at least 1. + s = s ? s : 1; + _typeTags.reserve(s); + _values.reserve(s); + } + +private: + std::vector<TypeTags> _typeTags; + std::vector<Value> _values; +}; + +/** + * This is a set of unique values with the same interface as Array. + */ +class ArraySet { + struct Hash { + size_t operator()(const std::pair<TypeTags, Value>& p) const { + return hashValue(p.first, p.second); + } + }; + struct Eq { + bool operator()(const std::pair<TypeTags, Value>& lhs, + const std::pair<TypeTags, Value>& rhs) const { + auto [tag, val] = compareValue(lhs.first, lhs.second, rhs.first, rhs.second); + + if (tag != TypeTags::NumberInt32 || val != 0) { + return false; + } else { + return true; + } + } + }; + using SetType = absl::flat_hash_set<std::pair<TypeTags, Value>, Hash, Eq>; + +public: + using iterator = SetType::iterator; + + ArraySet() = default; + ArraySet(const ArraySet& other) { + reserve(other._values.size()); + for (const auto& p : other._values) { + const auto copy = copyValue(p.first, p.second); + ValueGuard guard{copy.first, copy.second}; + _values.insert(copy); + guard.reset(); + } + } + ArraySet(ArraySet&&) = default; + ~ArraySet() { + for (const auto& p : _values) { + releaseValue(p.first, p.second); + } + } + + void push_back(TypeTags tag, Value val); + + auto& values() noexcept { + return _values; + } + + auto size() const noexcept { + return _values.size(); + } + void reserve(size_t s) { + // Normalize to at least 1. + s = s ? s : 1; + _values.reserve(s); + } + +private: + SetType _values; +}; + +constexpr size_t kSmallStringThreshold = 8; +using ObjectIdType = std::array<uint8_t, 12>; +static_assert(sizeof(ObjectIdType) == 12); + +inline char* getSmallStringView(Value& val) noexcept { + return reinterpret_cast<char*>(&val); +} + +inline char* getBigStringView(Value val) noexcept { + return reinterpret_cast<char*>(val); +} + +inline char* getRawPointerView(Value val) noexcept { + return reinterpret_cast<char*>(val); +} + +template <typename T> +T readFromMemory(const char* memory) noexcept { + T val; + memcpy(&val, memory, sizeof(T)); + return val; +} + +template <typename T> +T readFromMemory(const unsigned char* memory) noexcept { + T val; + memcpy(&val, memory, sizeof(T)); + return val; +} + +template <typename T> +size_t writeToMemory(unsigned char* memory, const T val) noexcept { + memcpy(memory, &val, sizeof(T)); + + return sizeof(T); +} + +template <typename T> +Value bitcastFrom(const T in) noexcept { + static_assert(sizeof(Value) >= sizeof(T)); + + // Casting from pointer to integer value is OK. + if constexpr (std::is_pointer_v<T>) { + return reinterpret_cast<Value>(in); + } + Value val{0}; + memcpy(&val, &in, sizeof(T)); + return val; +} + +template <typename T> +T bitcastTo(const Value in) noexcept { + // Casting from integer value to pointer is OK. + if constexpr (std::is_pointer_v<T>) { + static_assert(sizeof(Value) == sizeof(T)); + return reinterpret_cast<T>(in); + } else if constexpr (std::is_same_v<Decimal128, T>) { + static_assert(sizeof(Value) == sizeof(T*)); + T val; + memcpy(&val, getRawPointerView(in), sizeof(T)); + return val; + } else { + static_assert(sizeof(Value) >= sizeof(T)); + T val; + memcpy(&val, &in, sizeof(T)); + return val; + } +} + +inline std::string_view getStringView(TypeTags tag, Value& val) noexcept { + if (tag == TypeTags::StringSmall) { + return std::string_view(getSmallStringView(val)); + } else if (tag == TypeTags::StringBig) { + return std::string_view(getBigStringView(val)); + } else if (tag == TypeTags::bsonString) { + auto bsonstr = getRawPointerView(val); + return std::string_view(bsonstr + 4, + ConstDataView(bsonstr).read<LittleEndian<uint32_t>>() - 1); + } + MONGO_UNREACHABLE; +} + +inline std::pair<TypeTags, Value> makeNewString(std::string_view input) { + size_t len = input.size(); + if (len < kSmallStringThreshold - 1) { + Value smallString; + // This is OK - we are aliasing to char*. + auto stringAlias = getSmallStringView(smallString); + memcpy(stringAlias, input.data(), len); + stringAlias[len] = 0; + return {TypeTags::StringSmall, smallString}; + } else { + auto str = new char[len + 1]; + memcpy(str, input.data(), len); + str[len] = 0; + return {TypeTags::StringBig, reinterpret_cast<Value>(str)}; + } +} + +inline std::pair<TypeTags, Value> makeNewArray() { + auto a = new Array; + return {TypeTags::Array, reinterpret_cast<Value>(a)}; +} + +inline std::pair<TypeTags, Value> makeNewArraySet() { + auto a = new ArraySet; + return {TypeTags::ArraySet, reinterpret_cast<Value>(a)}; +} + +inline std::pair<TypeTags, Value> makeCopyArray(const Array& inA) { + auto a = new Array(inA); + return {TypeTags::Array, reinterpret_cast<Value>(a)}; +} + +inline std::pair<TypeTags, Value> makeCopyArraySet(const ArraySet& inA) { + auto a = new ArraySet(inA); + return {TypeTags::ArraySet, reinterpret_cast<Value>(a)}; +} + +inline Array* getArrayView(Value val) noexcept { + return reinterpret_cast<Array*>(val); +} + +inline ArraySet* getArraySetView(Value val) noexcept { + return reinterpret_cast<ArraySet*>(val); +} + +inline std::pair<TypeTags, Value> makeNewObject() { + auto o = new Object; + return {TypeTags::Object, reinterpret_cast<Value>(o)}; +} + +inline std::pair<TypeTags, Value> makeCopyObject(const Object& inO) { + auto o = new Object(inO); + return {TypeTags::Object, reinterpret_cast<Value>(o)}; +} + +inline Object* getObjectView(Value val) noexcept { + return reinterpret_cast<Object*>(val); +} + +inline std::pair<TypeTags, Value> makeNewObjectId() { + auto o = new ObjectIdType; + return {TypeTags::ObjectId, reinterpret_cast<Value>(o)}; +} + +inline std::pair<TypeTags, Value> makeCopyObjectId(const ObjectIdType& inO) { + auto o = new ObjectIdType(inO); + return {TypeTags::ObjectId, reinterpret_cast<Value>(o)}; +} + +inline ObjectIdType* getObjectIdView(Value val) noexcept { + return reinterpret_cast<ObjectIdType*>(val); +} + +inline Decimal128* getDecimalView(Value val) noexcept { + return reinterpret_cast<Decimal128*>(val); +} + +inline std::pair<TypeTags, Value> makeCopyDecimal(const Decimal128& inD) { + auto o = new Decimal128(inD); + return {TypeTags::NumberDecimal, reinterpret_cast<Value>(o)}; +} + +inline KeyString::Value* getKeyStringView(Value val) noexcept { + return reinterpret_cast<KeyString::Value*>(val); +} + +inline pcrecpp::RE* getPrceRegexView(Value val) noexcept { + return reinterpret_cast<pcrecpp::RE*>(val); +} + +std::pair<TypeTags, Value> makeCopyKeyString(const KeyString::Value& inKey); + +std::pair<TypeTags, Value> makeCopyPcreRegex(const pcrecpp::RE&); + +void releaseValue(TypeTags tag, Value val) noexcept; + +inline std::pair<TypeTags, Value> copyValue(TypeTags tag, Value val) { + switch (tag) { + case TypeTags::NumberDecimal: + return makeCopyDecimal(bitcastTo<Decimal128>(val)); + case TypeTags::Array: + return makeCopyArray(*getArrayView(val)); + case TypeTags::ArraySet: + return makeCopyArraySet(*getArraySetView(val)); + case TypeTags::Object: + return makeCopyObject(*getObjectView(val)); + case TypeTags::StringBig: { + auto src = getBigStringView(val); + auto len = strlen(src); + auto dst = new char[len + 1]; + memcpy(dst, src, len); + dst[len] = 0; + return {TypeTags::StringBig, bitcastFrom(dst)}; + } + case TypeTags::bsonString: { + auto bsonstr = getRawPointerView(val); + auto src = bsonstr + 4; + auto size = ConstDataView(bsonstr).read<LittleEndian<uint32_t>>(); + return makeNewString(std::string_view(src, size - 1)); + } + case TypeTags::ObjectId: { + return makeCopyObjectId(*getObjectIdView(val)); + } + case TypeTags::bsonObject: { + auto bson = getRawPointerView(val); + auto size = ConstDataView(bson).read<LittleEndian<uint32_t>>(); + auto dst = new uint8_t[size]; + memcpy(dst, bson, size); + return {TypeTags::bsonObject, bitcastFrom(dst)}; + } + case TypeTags::bsonObjectId: { + auto bson = getRawPointerView(val); + auto size = sizeof(ObjectIdType); + auto dst = new uint8_t[size]; + memcpy(dst, bson, size); + return {TypeTags::bsonObjectId, bitcastFrom(dst)}; + } + case TypeTags::bsonArray: { + auto bson = getRawPointerView(val); + auto size = ConstDataView(bson).read<LittleEndian<uint32_t>>(); + auto dst = new uint8_t[size]; + memcpy(dst, bson, size); + return {TypeTags::bsonArray, bitcastFrom(dst)}; + } + case TypeTags::ksValue: + return makeCopyKeyString(*getKeyStringView(val)); + case TypeTags::pcreRegex: + return makeCopyPcreRegex(*getPrceRegexView(val)); + default: + break; + } + + return {tag, val}; +} + +/** + * Implicit conversions of numerical types. + */ +template <typename T> +inline T numericCast(TypeTags tag, Value val) noexcept { + switch (tag) { + case TypeTags::NumberInt32: + if constexpr (std::is_same_v<T, Decimal128>) { + return Decimal128(bitcastTo<int32_t>(val)); + } else { + return bitcastTo<int32_t>(val); + } + case TypeTags::NumberInt64: + if constexpr (std::is_same_v<T, Decimal128>) { + return Decimal128(bitcastTo<int64_t>(val)); + } else { + return bitcastTo<int64_t>(val); + } + case TypeTags::NumberDouble: + if constexpr (std::is_same_v<T, Decimal128>) { + return Decimal128(bitcastTo<double>(val)); + } else { + return bitcastTo<double>(val); + } + case TypeTags::NumberDecimal: + if constexpr (std::is_same_v<T, Decimal128>) { + return bitcastTo<Decimal128>(val); + } + MONGO_UNREACHABLE; + default: + MONGO_UNREACHABLE; + } +} + +inline TypeTags getWidestNumericalType(TypeTags lhsTag, TypeTags rhsTag) noexcept { + if (lhsTag == TypeTags::NumberDecimal || rhsTag == TypeTags::NumberDecimal) { + return TypeTags::NumberDecimal; + } else if (lhsTag == TypeTags::NumberDouble || rhsTag == TypeTags::NumberDouble) { + return TypeTags::NumberDouble; + } else if (lhsTag == TypeTags::NumberInt64 || rhsTag == TypeTags::NumberInt64) { + return TypeTags::NumberInt64; + } else if (lhsTag == TypeTags::NumberInt32 || rhsTag == TypeTags::NumberInt32) { + return TypeTags::NumberInt32; + } else { + MONGO_UNREACHABLE; + } +} + +class SlotAccessor { +public: + virtual ~SlotAccessor() = 0; + /** + * Returns a non-owning view of value currently stored in the slot. The returned value is valid + * until the content of this slot changes (usually as a result of calling getNext()). If the + * caller needs to hold onto the value longer then it must make a copy of the value. + */ + virtual std::pair<TypeTags, Value> getViewOfValue() const = 0; + + /** + * Sometimes it may be determined that a caller is the last one to access this slot. If that is + * the case then the caller can use this optimized method to move out the value out of the slot + * saving the extra copy operation. Not all slots own the values stored in them so they must + * make a deep copy. The returned value is owned by the caller. + */ + virtual std::pair<TypeTags, Value> copyOrMoveValue() = 0; +}; +inline SlotAccessor::~SlotAccessor() {} + +class ViewOfValueAccessor final : public SlotAccessor { +public: + // Return non-owning view of the value. + std::pair<TypeTags, Value> getViewOfValue() const override { + return {_tag, _val}; + } + + // Return a copy. + std::pair<TypeTags, Value> copyOrMoveValue() override { + return copyValue(_tag, _val); + } + + void reset() { + reset(TypeTags::Nothing, 0); + } + + void reset(TypeTags tag, Value val) { + _tag = tag; + _val = val; + } + +private: + TypeTags _tag{TypeTags::Nothing}; + Value _val{0}; +}; + +class OwnedValueAccessor final : public SlotAccessor { +public: + OwnedValueAccessor() = default; + OwnedValueAccessor(const OwnedValueAccessor& other) { + if (other._owned) { + auto [tag, val] = copyValue(other._tag, other._val); + _tag = tag; + _val = val; + _owned = true; + } else { + _tag = other._tag; + _val = other._val; + _owned = false; + } + } + OwnedValueAccessor(OwnedValueAccessor&& other) noexcept { + _tag = other._tag; + _val = other._val; + _owned = other._owned; + + other._owned = false; + } + + ~OwnedValueAccessor() { + release(); + } + + // Copy and swap idiom for a single copy/move assignment operator. + OwnedValueAccessor& operator=(OwnedValueAccessor other) noexcept { + std::swap(_tag, other._tag); + std::swap(_val, other._val); + std::swap(_owned, other._owned); + return *this; + } + + // Return non-owning view of the value. + std::pair<TypeTags, Value> getViewOfValue() const override { + return {_tag, _val}; + } + + std::pair<TypeTags, Value> copyOrMoveValue() override { + if (_owned) { + _owned = false; + return {_tag, _val}; + } else { + return copyValue(_tag, _val); + } + } + + void reset() { + reset(TypeTags::Nothing, 0); + } + + void reset(TypeTags tag, Value val) { + reset(true, tag, val); + } + + void reset(bool owned, TypeTags tag, Value val) { + release(); + + _tag = tag; + _val = val; + _owned = owned; + } + +private: + void release() { + if (_owned) { + releaseValue(_tag, _val); + _owned = false; + } + } + + TypeTags _tag{TypeTags::Nothing}; + Value _val; + bool _owned{false}; +}; + +class ObjectEnumerator { +public: + ObjectEnumerator() = default; + ObjectEnumerator(TypeTags tag, Value val) { + reset(tag, val); + } + void reset(TypeTags tag, Value val) { + _tagObject = tag; + _valObject = val; + _object = nullptr; + _index = 0; + + if (tag == TypeTags::Object) { + _object = getObjectView(val); + } else if (tag == TypeTags::bsonObject) { + auto bson = getRawPointerView(val); + _objectCurrent = bson + 4; + _objectEnd = bson + ConstDataView(bson).read<LittleEndian<uint32_t>>(); + } else { + MONGO_UNREACHABLE; + } + } + std::pair<TypeTags, Value> getViewOfValue() const; + std::string_view getFieldName() const; + + bool atEnd() const { + if (_object) { + return _index == _object->size(); + } else { + return *_objectCurrent == 0; + } + } + + bool advance(); + +private: + TypeTags _tagObject{TypeTags::Nothing}; + Value _valObject{0}; + + // Object + Object* _object{nullptr}; + size_t _index{0}; + + // bsonObject + const char* _objectCurrent{nullptr}; + const char* _objectEnd{nullptr}; +}; + +class ArrayEnumerator { +public: + ArrayEnumerator() = default; + ArrayEnumerator(TypeTags tag, Value val) { + reset(tag, val); + } + void reset(TypeTags tag, Value val) { + _tagArray = tag; + _valArray = val; + _array = nullptr; + _arraySet = nullptr; + _index = 0; + + if (tag == TypeTags::Array) { + _array = getArrayView(val); + } else if (tag == TypeTags::ArraySet) { + _arraySet = getArraySetView(val); + _iter = _arraySet->values().begin(); + } else if (tag == TypeTags::bsonArray) { + auto bson = getRawPointerView(val); + _arrayCurrent = bson + 4; + _arrayEnd = bson + ConstDataView(bson).read<LittleEndian<uint32_t>>(); + } else { + MONGO_UNREACHABLE; + } + } + std::pair<TypeTags, Value> getViewOfValue() const; + + bool atEnd() const { + if (_array) { + return _index == _array->size(); + } else if (_arraySet) { + return _iter == _arraySet->values().end(); + } else { + return *_arrayCurrent == 0; + } + } + + bool advance(); + +private: + TypeTags _tagArray{TypeTags::Nothing}; + Value _valArray{0}; + + // Array + Array* _array{nullptr}; + size_t _index{0}; + + // ArraySet + ArraySet* _arraySet{nullptr}; + ArraySet::iterator _iter; + + // bsonArray + const char* _arrayCurrent{nullptr}; + const char* _arrayEnd{nullptr}; +}; + +class ArrayAccessor final : public SlotAccessor { +public: + void reset(TypeTags tag, Value val) { + _enumerator.reset(tag, val); + } + + // Return non-owning view of the value. + std::pair<TypeTags, Value> getViewOfValue() const override { + return _enumerator.getViewOfValue(); + } + std::pair<TypeTags, Value> copyOrMoveValue() override { + // We can never move out values from array. + auto [tag, val] = getViewOfValue(); + return copyValue(tag, val); + } + + bool atEnd() const { + return _enumerator.atEnd(); + } + + bool advance() { + return _enumerator.advance(); + } + +private: + ArrayEnumerator _enumerator; +}; + +struct MaterializedRow { + std::vector<OwnedValueAccessor> _fields; + + void makeOwned() { + for (auto& f : _fields) { + auto [tag, val] = f.getViewOfValue(); + auto [copyTag, copyVal] = copyValue(tag, val); + f.reset(copyTag, copyVal); + } + } + bool operator==(const MaterializedRow& rhs) const { + for (size_t idx = 0; idx < _fields.size(); ++idx) { + auto [lhsTag, lhsVal] = _fields[idx].getViewOfValue(); + auto [rhsTag, rhsVal] = rhs._fields[idx].getViewOfValue(); + auto [tag, val] = compareValue(lhsTag, lhsVal, rhsTag, rhsVal); + + if (tag != TypeTags::NumberInt32 || val != 0) { + return false; + } + } + + return true; + } +}; + +struct MaterializedRowComparator { + const std::vector<SortDirection>& _direction; + // TODO - add collator and whatnot. + + MaterializedRowComparator(const std::vector<value::SortDirection>& direction) + : _direction(direction) {} + + bool operator()(const MaterializedRow& lhs, const MaterializedRow& rhs) const { + for (size_t idx = 0; idx < lhs._fields.size(); ++idx) { + auto [lhsTag, lhsVal] = lhs._fields[idx].getViewOfValue(); + auto [rhsTag, rhsVal] = rhs._fields[idx].getViewOfValue(); + auto [tag, val] = compareValue(lhsTag, lhsVal, rhsTag, rhsVal); + if (tag != TypeTags::NumberInt32) { + return false; + } + if (bitcastTo<int32_t>(val) < 0 && _direction[idx] == SortDirection::Ascending) { + return true; + } + if (bitcastTo<int32_t>(val) > 0 && _direction[idx] == SortDirection::Descending) { + return true; + } + if (bitcastTo<int32_t>(val) != 0) { + return false; + } + } + + return false; + } +}; +struct MaterializedRowHasher { + std::size_t operator()(const MaterializedRow& k) const { + size_t res = 17; + for (auto& f : k._fields) { + auto [tag, val] = f.getViewOfValue(); + res = res * 31 + hashValue(tag, val); + } + return res; + } +}; + +template <typename T> +class MaterializedRowKeyAccessor final : public SlotAccessor { +public: + MaterializedRowKeyAccessor(T& it, size_t slot) : _it(it), _slot(slot) {} + + std::pair<TypeTags, Value> getViewOfValue() const override { + return _it->first._fields[_slot].getViewOfValue(); + } + std::pair<TypeTags, Value> copyOrMoveValue() override { + // We can never move out values from keys. + auto [tag, val] = getViewOfValue(); + return copyValue(tag, val); + } + +private: + T& _it; + size_t _slot; +}; + +template <typename T> +class MaterializedRowValueAccessor final : public SlotAccessor { +public: + MaterializedRowValueAccessor(T& it, size_t slot) : _it(it), _slot(slot) {} + + std::pair<TypeTags, Value> getViewOfValue() const override { + return _it->second._fields[_slot].getViewOfValue(); + } + std::pair<TypeTags, Value> copyOrMoveValue() override { + return _it->second._fields[_slot].copyOrMoveValue(); + } + + void reset(bool owned, TypeTags tag, Value val) { + _it->second._fields[_slot].reset(owned, tag, val); + } + +private: + T& _it; + size_t _slot; +}; + +template <typename T> +class MaterializedRowAccessor final : public SlotAccessor { +public: + MaterializedRowAccessor(T& container, const size_t& it, size_t slot) + : _container(container), _it(it), _slot(slot) {} + + std::pair<TypeTags, Value> getViewOfValue() const override { + return _container[_it]._fields[_slot].getViewOfValue(); + } + std::pair<TypeTags, Value> copyOrMoveValue() override { + return _container[_it]._fields[_slot].copyOrMoveValue(); + } + + void reset(bool owned, TypeTags tag, Value val) { + _container[_it]._fields[_slot].reset(owned, tag, val); + } + +private: + T& _container; + const size_t& _it; + const size_t _slot; +}; + +/** + * Commonly used containers + */ +template <typename T> +using SlotMap = absl::flat_hash_map<SlotId, T>; +using SlotAccessorMap = SlotMap<SlotAccessor*>; +using FieldAccessorMap = absl::flat_hash_map<std::string, std::unique_ptr<ViewOfValueAccessor>>; +using SlotSet = absl::flat_hash_set<SlotId>; +using SlotVector = std::vector<SlotId>; + +} // namespace value +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/vm/arith.cpp b/src/mongo/db/exec/sbe/vm/arith.cpp new file mode 100644 index 00000000000..98f4f110c2b --- /dev/null +++ b/src/mongo/db/exec/sbe/vm/arith.cpp @@ -0,0 +1,277 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/db/exec/sbe/vm/vm.h" + +namespace mongo { +namespace sbe { +namespace vm { + +using namespace value; + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericAdd(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue) { + if (value::isNumber(lhsTag) && value::isNumber(rhsTag)) { + switch (getWidestNumericalType(lhsTag, rhsTag)) { + case value::TypeTags::NumberInt32: { + auto result = + numericCast<int32_t>(lhsTag, lhsValue) + numericCast<int32_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt32, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberInt64: { + auto result = + numericCast<int64_t>(lhsTag, lhsValue) + numericCast<int64_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt64, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDouble: { + auto result = + numericCast<double>(lhsTag, lhsValue) + numericCast<double>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberDouble, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDecimal: { + auto result = numericCast<Decimal128>(lhsTag, lhsValue) + .add(numericCast<Decimal128>(rhsTag, rhsValue)); + auto [tag, val] = value::makeCopyDecimal(result); + return {true, tag, val}; + } + default: + MONGO_UNREACHABLE; + } + } + + return {false, value::TypeTags::Nothing, 0}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericSub(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue) { + if (value::isNumber(lhsTag) && value::isNumber(rhsTag)) { + switch (getWidestNumericalType(lhsTag, rhsTag)) { + case value::TypeTags::NumberInt32: { + auto result = + numericCast<int32_t>(lhsTag, lhsValue) - numericCast<int32_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt32, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberInt64: { + auto result = + numericCast<int64_t>(lhsTag, lhsValue) - numericCast<int64_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt64, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDouble: { + auto result = + numericCast<double>(lhsTag, lhsValue) - numericCast<double>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberDouble, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDecimal: { + auto result = numericCast<Decimal128>(lhsTag, lhsValue) + .subtract(numericCast<Decimal128>(rhsTag, rhsValue)); + auto [tag, val] = value::makeCopyDecimal(result); + return {true, tag, val}; + } + default: + MONGO_UNREACHABLE; + } + } + + return {false, value::TypeTags::Nothing, 0}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericMul(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue) { + if (value::isNumber(lhsTag) && value::isNumber(rhsTag)) { + switch (getWidestNumericalType(lhsTag, rhsTag)) { + case value::TypeTags::NumberInt32: { + auto result = + numericCast<int32_t>(lhsTag, lhsValue) * numericCast<int32_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt32, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberInt64: { + auto result = + numericCast<int64_t>(lhsTag, lhsValue) * numericCast<int64_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt64, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDouble: { + auto result = + numericCast<double>(lhsTag, lhsValue) * numericCast<double>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberDouble, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDecimal: { + auto result = numericCast<Decimal128>(lhsTag, lhsValue) + .multiply(numericCast<Decimal128>(rhsTag, rhsValue)); + auto [tag, val] = value::makeCopyDecimal(result); + return {true, tag, val}; + } + default: + MONGO_UNREACHABLE; + } + } + + return {false, value::TypeTags::Nothing, 0}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericDiv(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue) { + if (value::isNumber(lhsTag) && value::isNumber(rhsTag)) { + switch (getWidestNumericalType(lhsTag, rhsTag)) { + case value::TypeTags::NumberInt32: { + auto result = + numericCast<int32_t>(lhsTag, lhsValue) / numericCast<int32_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt32, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberInt64: { + auto result = + numericCast<int64_t>(lhsTag, lhsValue) / numericCast<int64_t>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberInt64, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDouble: { + auto result = + numericCast<double>(lhsTag, lhsValue) / numericCast<double>(rhsTag, rhsValue); + return {false, value::TypeTags::NumberDouble, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDecimal: { + auto result = numericCast<Decimal128>(lhsTag, lhsValue) + .divide(numericCast<Decimal128>(rhsTag, rhsValue)); + auto [tag, val] = value::makeCopyDecimal(result); + return {true, tag, val}; + } + default: + MONGO_UNREACHABLE; + } + } + + return {false, value::TypeTags::Nothing, 0}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericAbs(value::TypeTags operandTag, + value::Value operandValue) { + switch (operandTag) { + case value::TypeTags::NumberInt32: { + auto operand = value::bitcastTo<int32_t>(operandValue); + if (operand == std::numeric_limits<int32_t>::min()) { + return {false, value::TypeTags::NumberInt64, value::bitcastFrom(int64_t{operand})}; + } + + return {false, + value::TypeTags::NumberInt32, + value::bitcastFrom(operand >= 0 ? operand : -operand)}; + } + case value::TypeTags::NumberInt64: { + auto operand = value::bitcastTo<int64_t>(operandValue); + uassert(/* Intentionally duplicated */ 28680, + "can't take $abs of long long min", + operand != std::numeric_limits<int64_t>::min()); + return {false, + value::TypeTags::NumberInt64, + value::bitcastFrom(operand >= 0 ? operand : -operand)}; + } + case value::TypeTags::NumberDouble: { + auto operand = value::bitcastTo<double>(operandValue); + return {false, + value::TypeTags::NumberDouble, + value::bitcastFrom(operand >= 0 ? operand : -operand)}; + } + case value::TypeTags::NumberDecimal: { + auto operand = value::bitcastTo<Decimal128>(operandValue); + auto [tag, value] = value::makeCopyDecimal(operand.toAbs()); + return {true, tag, value}; + } + default: + return {false, value::TypeTags::Nothing, 0}; + } +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericNot(value::TypeTags tag, + value::Value value) { + if (tag == value::TypeTags::Boolean) { + return { + false, value::TypeTags::Boolean, value::bitcastFrom(!value::bitcastTo<bool>(value))}; + } else { + return {false, value::TypeTags::Nothing, 0}; + } +} + +std::pair<value::TypeTags, value::Value> ByteCode::genericCompareEq(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue) { + if ((value::isNumber(lhsTag) && value::isNumber(rhsTag)) || + (lhsTag == value::TypeTags::Date && rhsTag == value::TypeTags::Date) || + (lhsTag == value::TypeTags::Timestamp && rhsTag == value::TypeTags::Timestamp)) { + return genericNumericCompare(lhsTag, lhsValue, rhsTag, rhsValue, std::equal_to<>{}); + } else if (value::isString(lhsTag) && value::isString(rhsTag)) { + auto lhsStr = value::getStringView(lhsTag, lhsValue); + auto rhsStr = value::getStringView(rhsTag, rhsValue); + + return {value::TypeTags::Boolean, lhsStr.compare(rhsStr) == 0}; + } else if (lhsTag == value::TypeTags::Boolean && rhsTag == value::TypeTags::Boolean) { + return {value::TypeTags::Boolean, (lhsValue != 0) == (rhsValue != 0)}; + } else if (lhsTag == value::TypeTags::Null && rhsTag == value::TypeTags::Null) { + // This is where Mongo differs from SQL. + return {value::TypeTags::Boolean, true}; + } else if (lhsTag == value::TypeTags::ObjectId && rhsTag == value::TypeTags::ObjectId) { + return {value::TypeTags::Boolean, + (*value::getObjectIdView(lhsValue)) == (*value::getObjectIdView(rhsValue))}; + } else { + return {value::TypeTags::Nothing, 0}; + } +} + +std::pair<value::TypeTags, value::Value> ByteCode::genericCompareNeq(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue) { + auto [tag, val] = genericCompareEq(lhsTag, lhsValue, rhsTag, rhsValue); + if (tag == value::TypeTags::Boolean) { + return {tag, value::bitcastFrom(!value::bitcastTo<bool>(val))}; + } else { + return {tag, val}; + } +} + +std::pair<value::TypeTags, value::Value> ByteCode::compare3way(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue) { + if (lhsTag == value::TypeTags::Nothing || rhsTag == value::TypeTags::Nothing) { + return {value::TypeTags::Nothing, 0}; + } + + return value::compareValue(lhsTag, lhsValue, rhsTag, rhsValue); +} + +} // namespace vm +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp new file mode 100644 index 00000000000..29a93219b15 --- /dev/null +++ b/src/mongo/db/exec/sbe/vm/vm.cpp @@ -0,0 +1,1383 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/vm/vm.h" + +#include <pcrecpp.h> + +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/exec/sbe/values/value.h" +#include "mongo/db/storage/key_string.h" +#include "mongo/util/fail_point.h" + +MONGO_FAIL_POINT_DEFINE(failOnPoisonedFieldLookup); + +namespace mongo { +namespace sbe { +namespace vm { + +/* + * This table must be kept in sync with Instruction::Tags. It encodes how the instruction affects + * the stack; i.e. push(+1), pop(-1), or no effect. + */ +int Instruction::stackOffset[Instruction::Tags::lastInstruction] = { + 1, // pushConstVal + 1, // pushAccessVal + 1, // pushMoveVal + 1, // pushLocalVal + -1, // pop + 0, // swap + + -1, // add + -1, // sub + -1, // mul + -1, // div + 0, // negate + + 0, // logicNot + + -1, // less + -1, // lessEq + -1, // greater + -1, // greaterEq + -1, // eq + -1, // neq + -1, // cmp3w + + -1, // fillEmpty + + -1, // getField + + -1, // sum + -1, // min + -1, // max + -1, // first + -1, // last + + 0, // exists + 0, // isNull + 0, // isObject + 0, // isArray + 0, // isString + 0, // isNumber + + 0, // function is special, the stack offset is encoded in the instruction itself + + 0, // jmp + -1, // jmpTrue + 0, // jmpNothing + + -1, // fail +}; + +void CodeFragment::adjustStackSimple(const Instruction& i) { + _stackSize += Instruction::stackOffset[i.tag]; +} + +void CodeFragment::fixup(int offset) { + for (auto fixUp : _fixUps) { + auto ptr = instrs().data() + fixUp.offset; + int newOffset = value::readFromMemory<int>(ptr) + offset; + value::writeToMemory(ptr, newOffset); + } +} + +void CodeFragment::removeFixup(FrameId frameId) { + _fixUps.erase(std::remove_if(_fixUps.begin(), + _fixUps.end(), + [frameId](const auto& f) { return f.frameId == frameId; }), + _fixUps.end()); +} + +void CodeFragment::copyCodeAndFixup(const CodeFragment& from) { + for (auto fixUp : from._fixUps) { + fixUp.offset += _instrs.size(); + _fixUps.push_back(fixUp); + } + + _instrs.insert(_instrs.end(), from._instrs.begin(), from._instrs.end()); +} + +void CodeFragment::append(std::unique_ptr<CodeFragment> code) { + // Fixup before copying. + code->fixup(_stackSize); + + copyCodeAndFixup(*code); + + _stackSize += code->_stackSize; +} + +void CodeFragment::append(std::unique_ptr<CodeFragment> lhs, std::unique_ptr<CodeFragment> rhs) { + invariant(lhs->stackSize() == rhs->stackSize()); + + // Fixup before copying. + lhs->fixup(_stackSize); + rhs->fixup(_stackSize); + + copyCodeAndFixup(*lhs); + copyCodeAndFixup(*rhs); + + _stackSize += lhs->_stackSize; +} + +void CodeFragment::appendConstVal(value::TypeTags tag, value::Value val) { + Instruction i; + i.tag = Instruction::pushConstVal; + adjustStackSimple(i); + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(tag) + sizeof(val)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, tag); + offset += value::writeToMemory(offset, val); +} + +void CodeFragment::appendAccessVal(value::SlotAccessor* accessor) { + Instruction i; + i.tag = Instruction::pushAccessVal; + adjustStackSimple(i); + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(accessor)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, accessor); +} + +void CodeFragment::appendMoveVal(value::SlotAccessor* accessor) { + Instruction i; + i.tag = Instruction::pushMoveVal; + adjustStackSimple(i); + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(accessor)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, accessor); +} + +void CodeFragment::appendLocalVal(FrameId frameId, int stackOffset) { + Instruction i; + i.tag = Instruction::pushLocalVal; + adjustStackSimple(i); + + auto fixUpOffset = _instrs.size() + sizeof(Instruction); + _fixUps.push_back(FixUp{frameId, fixUpOffset}); + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(stackOffset)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, stackOffset); +} + +void CodeFragment::appendAdd() { + appendSimpleInstruction(Instruction::add); +} + +void CodeFragment::appendSub() { + appendSimpleInstruction(Instruction::sub); +} + +void CodeFragment::appendMul() { + appendSimpleInstruction(Instruction::mul); +} + +void CodeFragment::appendDiv() { + appendSimpleInstruction(Instruction::div); +} + +void CodeFragment::appendNegate() { + appendSimpleInstruction(Instruction::negate); +} + +void CodeFragment::appendNot() { + appendSimpleInstruction(Instruction::logicNot); +} + +void CodeFragment::appendSimpleInstruction(Instruction::Tags tag) { + Instruction i; + i.tag = tag; + adjustStackSimple(i); + + auto offset = allocateSpace(sizeof(Instruction)); + + offset += value::writeToMemory(offset, i); +} + +void CodeFragment::appendGetField() { + appendSimpleInstruction(Instruction::getField); +} + +void CodeFragment::appendSum() { + appendSimpleInstruction(Instruction::aggSum); +} + +void CodeFragment::appendMin() { + appendSimpleInstruction(Instruction::aggMin); +} + +void CodeFragment::appendMax() { + appendSimpleInstruction(Instruction::aggMax); +} + +void CodeFragment::appendFirst() { + appendSimpleInstruction(Instruction::aggFirst); +} + +void CodeFragment::appendLast() { + appendSimpleInstruction(Instruction::aggLast); +} + +void CodeFragment::appendExists() { + appendSimpleInstruction(Instruction::exists); +} + +void CodeFragment::appendIsNull() { + appendSimpleInstruction(Instruction::isNull); +} + +void CodeFragment::appendIsObject() { + appendSimpleInstruction(Instruction::isObject); +} + +void CodeFragment::appendIsArray() { + appendSimpleInstruction(Instruction::isArray); +} + +void CodeFragment::appendIsString() { + appendSimpleInstruction(Instruction::isString); +} + +void CodeFragment::appendIsNumber() { + appendSimpleInstruction(Instruction::isNumber); +} + +void CodeFragment::appendFunction(Builtin f, uint8_t arity) { + Instruction i; + i.tag = Instruction::function; + + // Account for consumed arguments + _stackSize -= arity; + // and the return value. + _stackSize += 1; + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(f) + sizeof(arity)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, f); + offset += value::writeToMemory(offset, arity); +} + +void CodeFragment::appendJump(int jumpOffset) { + Instruction i; + i.tag = Instruction::jmp; + adjustStackSimple(i); + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(jumpOffset)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, jumpOffset); +} + +void CodeFragment::appendJumpTrue(int jumpOffset) { + Instruction i; + i.tag = Instruction::jmpTrue; + adjustStackSimple(i); + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(jumpOffset)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, jumpOffset); +} + +void CodeFragment::appendJumpNothing(int jumpOffset) { + Instruction i; + i.tag = Instruction::jmpNothing; + adjustStackSimple(i); + + auto offset = allocateSpace(sizeof(Instruction) + sizeof(jumpOffset)); + + offset += value::writeToMemory(offset, i); + offset += value::writeToMemory(offset, jumpOffset); +} + +ByteCode::~ByteCode() { + auto size = _argStackOwned.size(); + invariant(_argStackTags.size() == size); + invariant(_argStackVals.size() == size); + for (size_t i = 0; i < size; ++i) { + if (_argStackOwned[i]) { + value::releaseValue(_argStackTags[i], _argStackVals[i]); + } + } +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::getField(value::TypeTags objTag, + value::Value objValue, + value::TypeTags fieldTag, + value::Value fieldValue) { + if (!value::isString(fieldTag)) { + return {false, value::TypeTags::Nothing, 0}; + } + + auto fieldStr = value::getStringView(fieldTag, fieldValue); + + if (MONGO_unlikely(failOnPoisonedFieldLookup.shouldFail())) { + uassert(4623399, "Lookup of $POISON", fieldStr != "POISON"); + } + + if (objTag == value::TypeTags::Object) { + auto [tag, val] = value::getObjectView(objValue)->getField(fieldStr); + return {false, tag, val}; + } else if (objTag == value::TypeTags::bsonObject) { + auto be = value::bitcastTo<const char*>(objValue); + auto end = be + ConstDataView(be).read<LittleEndian<uint32_t>>(); + ; + // Skip document length. + be += 4; + while (*be != 0) { + auto sv = bson::fieldNameView(be); + + if (sv == fieldStr) { + auto [tag, val] = bson::convertFrom(true, be, end, sv.size()); + return {false, tag, val}; + } + + be = bson::advance(be, sv.size()); + } + } + return {false, value::TypeTags::Nothing, 0}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggSum(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue) { + // Skip aggregation step if we don't have the input. + if (fieldTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; + } + + // Initialize the accumulator. + if (accTag == value::TypeTags::Nothing) { + accTag = value::TypeTags::NumberInt64; + accValue = 0; + } + + return genericAdd(accTag, accValue, fieldTag, fieldValue); +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMin(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue) { + // Skip aggregation step if we don't have the input. + if (fieldTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; + } + + // Initialize the accumulator. + if (accTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(fieldTag, fieldValue); + return {true, tag, val}; + } + + auto [tag, val] = genericCompare<std::less<>>(accTag, accValue, fieldTag, fieldValue); + if (tag == value::TypeTags::Boolean && val) { + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; + } else { + auto [tag, val] = value::copyValue(fieldTag, fieldValue); + return {true, tag, val}; + } +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggMax(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue) { + // Skip aggregation step if we don't have the input. + if (fieldTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; + } + + // Initialize the accumulator. + if (accTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(fieldTag, fieldValue); + return {true, tag, val}; + } + + auto [tag, val] = genericCompare<std::greater<>>(accTag, accValue, fieldTag, fieldValue); + if (tag == value::TypeTags::Boolean && val) { + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; + } else { + auto [tag, val] = value::copyValue(fieldTag, fieldValue); + return {true, tag, val}; + } +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggFirst(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue) { + // Skip aggregation step if we don't have the input. + if (fieldTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; + } + + // Initialize the accumulator. + if (accTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(fieldTag, fieldValue); + return {true, tag, val}; + } + + // Disregard the next value, always return the first one. + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggLast(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue) { + // Skip aggregation step if we don't have the input. + if (fieldTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(accTag, accValue); + return {true, tag, val}; + } + + // Initialize the accumulator. + if (accTag == value::TypeTags::Nothing) { + auto [tag, val] = value::copyValue(fieldTag, fieldValue); + return {true, tag, val}; + } + + // Disregard the accumulator, always return the next value. + auto [tag, val] = value::copyValue(fieldTag, fieldValue); + return {true, tag, val}; +} + + +bool hasSeparatorAt(size_t idx, std::string_view input, std::string_view separator) { + if (separator.size() + idx > input.size()) { + return false; + } + + return input.compare(idx, separator.size(), separator) == 0; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinSplit(uint8_t arity) { + auto [ownedSeparator, tagSeparator, valSeparator] = getFromStack(1); + auto [ownedInput, tagInput, valInput] = getFromStack(0); + + if (!value::isString(tagSeparator) || !value::isString(tagInput)) { + return {false, value::TypeTags::Nothing, 0}; + } + + auto input = value::getStringView(tagInput, valInput); + auto separator = value::getStringView(tagSeparator, valSeparator); + + auto [tag, val] = value::makeNewArray(); + auto arr = value::getArrayView(val); + value::ValueGuard guard{tag, val}; + + size_t splitStart = 0; + size_t splitPos; + while ((splitPos = input.find(separator, splitStart)) != std::string_view::npos) { + auto [tag, val] = value::makeNewString(input.substr(splitStart, splitPos - splitStart)); + arr->push_back(tag, val); + + splitPos += separator.size(); + splitStart = splitPos; + } + + // This is the last string. + { + auto [tag, val] = value::makeNewString(input.substr(splitStart, input.size() - splitStart)); + arr->push_back(tag, val); + } + + guard.reset(); + return {true, tag, val}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinDropFields(uint8_t arity) { + auto [ownedSeparator, tagInObj, valInObj] = getFromStack(0); + + // We operate only on objects. + if (!value::isObject(tagInObj)) { + return {false, value::TypeTags::Nothing, 0}; + } + + // Build the set of fields to drop. + std::set<std::string, std::less<>> restrictFieldsSet; + for (uint8_t idx = 1; idx < arity; ++idx) { + auto [owned, tag, val] = getFromStack(idx); + + if (!value::isString(tag)) { + return {false, value::TypeTags::Nothing, 0}; + } + + restrictFieldsSet.emplace(value::getStringView(tag, val)); + } + + auto [tag, val] = value::makeNewObject(); + auto obj = value::getObjectView(val); + value::ValueGuard guard{tag, val}; + + if (tagInObj == value::TypeTags::bsonObject) { + auto be = value::bitcastTo<const char*>(valInObj); + auto end = be + ConstDataView(be).read<LittleEndian<uint32_t>>(); + ; + // Skip document length. + be += 4; + while (*be != 0) { + auto sv = bson::fieldNameView(be); + + if (restrictFieldsSet.count(sv) == 0) { + auto [tag, val] = bson::convertFrom(false, be, end, sv.size()); + obj->push_back(sv, tag, val); + } + + be = bson::advance(be, sv.size()); + } + } else if (tagInObj == value::TypeTags::Object) { + auto objRoot = value::getObjectView(valInObj); + for (size_t idx = 0; idx < objRoot->size(); ++idx) { + std::string_view sv(objRoot->field(idx)); + + if (restrictFieldsSet.count(sv) == 0) { + + auto [tag, val] = objRoot->getAt(idx); + auto [copyTag, copyVal] = value::copyValue(tag, val); + obj->push_back(sv, copyTag, copyVal); + } + } + } + + guard.reset(); + return {true, tag, val}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinNewObj(uint8_t arity) { + std::vector<value::TypeTags> typeTags; + std::vector<value::Value> values; + std::vector<std::string> names; + + for (uint8_t idx = 0; idx < arity; idx += 2) { + { + auto [owned, tag, val] = getFromStack(idx); + + if (!value::isString(tag)) { + return {false, value::TypeTags::Nothing, 0}; + } + + names.emplace_back(value::getStringView(tag, val)); + } + { + auto [owned, tag, val] = getFromStack(idx + 1); + typeTags.push_back(tag); + values.push_back(val); + } + } + + auto [tag, val] = value::makeNewObject(); + auto obj = value::getObjectView(val); + value::ValueGuard guard{tag, val}; + + for (size_t idx = 0; idx < typeTags.size(); ++idx) { + auto [tagCopy, valCopy] = value::copyValue(typeTags[idx], values[idx]); + obj->push_back(names[idx], tagCopy, valCopy); + } + + guard.reset(); + return {true, tag, val}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinKeyStringToString(uint8_t arity) { + auto [owned, tagInKey, valInKey] = getFromStack(0); + + // We operate only on keys. + if (tagInKey != value::TypeTags::ksValue) { + return {false, value::TypeTags::Nothing, 0}; + } + + auto key = value::getKeyStringView(valInKey); + + auto [tagStr, valStr] = value::makeNewString(key->toString()); + + return {true, tagStr, valStr}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinNewKeyString(uint8_t arity) { + auto [_, tagInVersion, valInVersion] = getFromStack(0); + + if (!value::isNumber(tagInVersion) || + !(value::numericCast<int64_t>(tagInVersion, valInVersion) == 0 || + value::numericCast<int64_t>(tagInVersion, valInVersion) == 1)) { + return {false, value::TypeTags::Nothing, 0}; + } + KeyString::Version version = + static_cast<KeyString::Version>(value::numericCast<int64_t>(tagInVersion, valInVersion)); + + auto [__, tagInOrdering, valInOrdering] = getFromStack(1); + if (!value::isNumber(tagInOrdering)) { + return {false, value::TypeTags::Nothing, 0}; + } + auto orderingBits = value::numericCast<int32_t>(tagInOrdering, valInOrdering); + BSONObjBuilder bb; + for (size_t i = 0; i < Ordering::kMaxCompoundIndexKeys; ++i) { + bb.append(""_sd, (orderingBits & (1 << i)) ? 1 : 0); + } + + KeyString::HeapBuilder kb{version, Ordering::make(bb.done())}; + + for (size_t idx = 2; idx < arity - 1u; ++idx) { + auto [_, tag, val] = getFromStack(idx); + if (value::isNumber(tag)) { + auto num = value::numericCast<int64_t>(tag, val); + kb.appendNumberLong(num); + } else if (value::isString(tag)) { + auto str = value::getStringView(tag, val); + kb.appendString(StringData{str.data(), str.length()}); + } else { + uasserted(4822802, "unsuppored key string type"); + } + } + + auto [___, tagInDisrim, valInDiscrim] = getFromStack(arity - 1); + if (!value::isNumber(tagInDisrim)) { + return {false, value::TypeTags::Nothing, 0}; + } + auto discrimNum = value::numericCast<int64_t>(tagInDisrim, valInDiscrim); + if (discrimNum < 0 || discrimNum > 2) { + return {false, value::TypeTags::Nothing, 0}; + } + + kb.appendDiscriminator(static_cast<KeyString::Discriminator>(discrimNum)); + + return {true, value::TypeTags::ksValue, value::bitcastFrom(new KeyString::Value(kb.release()))}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAbs(uint8_t arity) { + invariant(arity == 1); + + auto [_, tagOperand, valOperand] = getFromStack(0); + + return genericAbs(tagOperand, valOperand); +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAddToArray(uint8_t arity) { + auto [ownAgg, tagAgg, valAgg] = getFromStack(0); + auto [_, tagField, valField] = getFromStack(1); + + // Create a new array is it does not exist yet. + if (tagAgg == value::TypeTags::Nothing) { + auto [tagNewAgg, valNewAgg] = value::makeNewArray(); + ownAgg = true; + tagAgg = tagNewAgg; + valAgg = valNewAgg; + } else { + // Take ownership of the accumulator. + topStack(false, value::TypeTags::Nothing, 0); + } + value::ValueGuard guard{tagAgg, valAgg}; + + invariant(ownAgg && tagAgg == value::TypeTags::Array); + auto arr = value::getArrayView(valAgg); + + // And push back the value. Note that array will ignore Nothing. + auto [tagCopy, valCopy] = value::copyValue(tagField, valField); + arr->push_back(tagCopy, valCopy); + + guard.reset(); + return {ownAgg, tagAgg, valAgg}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAddToSet(uint8_t arity) { + auto [ownAgg, tagAgg, valAgg] = getFromStack(0); + auto [_, tagField, valField] = getFromStack(1); + + // Create a new array is it does not exist yet. + if (tagAgg == value::TypeTags::Nothing) { + auto [tagNewAgg, valNewAgg] = value::makeNewArraySet(); + ownAgg = true; + tagAgg = tagNewAgg; + valAgg = valNewAgg; + } else { + // Take ownership of the accumulator + topStack(false, value::TypeTags::Nothing, 0); + } + value::ValueGuard guard{tagAgg, valAgg}; + + invariant(ownAgg && tagAgg == value::TypeTags::ArraySet); + auto arr = value::getArraySetView(valAgg); + + // And push back the value. Note that array will ignore Nothing. + auto [tagCopy, valCopy] = value::copyValue(tagField, valField); + arr->push_back(tagCopy, valCopy); + + guard.reset(); + return {ownAgg, tagAgg, valAgg}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinRegexMatch(uint8_t arity) { + invariant(arity == 2); + + auto [ownedPcreRegex, typeTagPcreRegex, valuePcreRegex] = getFromStack(0); + auto [ownedInputStr, typeTagInputStr, valueInputStr] = getFromStack(1); + + if (!value::isString(typeTagInputStr) || typeTagPcreRegex != value::TypeTags::pcreRegex) { + return {false, value::TypeTags::Nothing, 0}; + } + + auto stringView = value::getStringView(typeTagInputStr, valueInputStr); + pcrecpp::StringPiece pcreStringView{stringView.data(), static_cast<int>(stringView.size())}; + + auto pcreRegex = value::getPrceRegexView(valuePcreRegex); + auto regexMatchResult = pcreRegex->PartialMatch(pcreStringView); + + return {false, value::TypeTags::Boolean, regexMatchResult}; +} + +std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builtin f, + uint8_t arity) { + switch (f) { + case Builtin::split: + return builtinSplit(arity); + case Builtin::regexMatch: + return builtinRegexMatch(arity); + case Builtin::dropFields: + return builtinDropFields(arity); + case Builtin::newObj: + return builtinNewObj(arity); + case Builtin::ksToString: + return builtinKeyStringToString(arity); + case Builtin::newKs: + return builtinNewKeyString(arity); + case Builtin::abs: + return builtinAbs(arity); + case Builtin::addToArray: + return builtinAddToArray(arity); + case Builtin::addToSet: + return builtinAddToSet(arity); + } + + MONGO_UNREACHABLE; +} + +std::tuple<uint8_t, value::TypeTags, value::Value> ByteCode::run(CodeFragment* code) { + auto pcPointer = code->instrs().data(); + auto pcEnd = pcPointer + code->instrs().size(); + + for (;;) { + if (pcPointer == pcEnd) { + break; + } else { + Instruction i = value::readFromMemory<Instruction>(pcPointer); + pcPointer += sizeof(i); + switch (i.tag) { + case Instruction::pushConstVal: { + auto tag = value::readFromMemory<value::TypeTags>(pcPointer); + pcPointer += sizeof(tag); + auto val = value::readFromMemory<value::Value>(pcPointer); + pcPointer += sizeof(val); + + pushStack(false, tag, val); + + break; + } + case Instruction::pushAccessVal: { + auto accessor = value::readFromMemory<value::SlotAccessor*>(pcPointer); + pcPointer += sizeof(accessor); + + auto [tag, val] = accessor->getViewOfValue(); + pushStack(false, tag, val); + + break; + } + case Instruction::pushMoveVal: { + auto accessor = value::readFromMemory<value::SlotAccessor*>(pcPointer); + pcPointer += sizeof(accessor); + + auto [tag, val] = accessor->copyOrMoveValue(); + pushStack(true, tag, val); + + break; + } + case Instruction::pushLocalVal: { + auto stackOffset = value::readFromMemory<int>(pcPointer); + pcPointer += sizeof(stackOffset); + + auto [owned, tag, val] = getFromStack(stackOffset); + + pushStack(false, tag, val); + + break; + } + case Instruction::pop: { + auto [owned, tag, val] = getFromStack(0); + popStack(); + + if (owned) { + value::releaseValue(tag, val); + } + + break; + } + case Instruction::swap: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(1); + + // Swap values only if they are not physically same. + // Note - this has huge consequences for the memory management, it allows to + // return owned values from the let expressions. + if (!(rhsTag == lhsTag && rhsVal == lhsVal)) { + setStack(0, lhsOwned, lhsTag, lhsVal); + setStack(1, rhsOwned, rhsTag, rhsVal); + } else { + // The values are physically same then the top of the stack must never ever + // be owned. + invariant(!rhsOwned); + } + + break; + } + case Instruction::add: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = genericAdd(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::sub: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = genericSub(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::mul: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = genericMul(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::div: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = genericDiv(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::negate: { + auto [owned, tag, val] = getFromStack(0); + + auto [resultOwned, resultTag, resultVal] = + genericSub(value::TypeTags::NumberInt32, 0, tag, val); + + topStack(resultOwned, resultTag, resultVal); + + if (owned) { + value::releaseValue(resultTag, resultVal); + } + + break; + } + case Instruction::logicNot: { + auto [owned, tag, val] = getFromStack(0); + + auto [resultOwned, resultTag, resultVal] = genericNot(tag, val); + + topStack(resultOwned, resultTag, resultVal); + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::less: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [tag, val] = genericCompare<std::less<>>(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(false, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::lessEq: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [tag, val] = + genericCompare<std::less_equal<>>(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(false, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::greater: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [tag, val] = + genericCompare<std::greater<>>(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(false, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::greaterEq: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [tag, val] = + genericCompare<std::greater_equal<>>(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(false, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::eq: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [tag, val] = genericCompareEq(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(false, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::neq: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [tag, val] = genericCompareNeq(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(false, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::cmp3w: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [tag, val] = compare3way(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(false, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::fillEmpty: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + if (lhsTag == value::TypeTags::Nothing) { + topStack(rhsOwned, rhsTag, rhsVal); + + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + } else { + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + } + break; + } + case Instruction::getField: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = getField(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::aggSum: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = aggSum(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::aggMin: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = aggMin(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::aggMax: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = aggMax(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::aggFirst: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = aggFirst(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::aggLast: { + auto [rhsOwned, rhsTag, rhsVal] = getFromStack(0); + popStack(); + auto [lhsOwned, lhsTag, lhsVal] = getFromStack(0); + + auto [owned, tag, val] = aggLast(lhsTag, lhsVal, rhsTag, rhsVal); + + topStack(owned, tag, val); + + if (rhsOwned) { + value::releaseValue(rhsTag, rhsVal); + } + if (lhsOwned) { + value::releaseValue(lhsTag, lhsVal); + } + break; + } + case Instruction::exists: { + auto [owned, tag, val] = getFromStack(0); + + topStack(false, value::TypeTags::Boolean, tag != value::TypeTags::Nothing); + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::isNull: { + auto [owned, tag, val] = getFromStack(0); + + if (tag != value::TypeTags::Nothing) { + topStack(false, value::TypeTags::Boolean, tag == value::TypeTags::Null); + } + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::isObject: { + auto [owned, tag, val] = getFromStack(0); + + if (tag != value::TypeTags::Nothing) { + topStack(false, value::TypeTags::Boolean, value::isObject(tag)); + } + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::isArray: { + auto [owned, tag, val] = getFromStack(0); + + if (tag != value::TypeTags::Nothing) { + topStack(false, value::TypeTags::Boolean, value::isArray(tag)); + } + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::isString: { + auto [owned, tag, val] = getFromStack(0); + + if (tag != value::TypeTags::Nothing) { + topStack(false, value::TypeTags::Boolean, value::isString(tag)); + } + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::isNumber: { + auto [owned, tag, val] = getFromStack(0); + + if (tag != value::TypeTags::Nothing) { + topStack(false, value::TypeTags::Boolean, value::isNumber(tag)); + } + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::function: { + auto f = value::readFromMemory<Builtin>(pcPointer); + pcPointer += sizeof(f); + auto arity = value::readFromMemory<uint8_t>(pcPointer); + pcPointer += sizeof(arity); + + auto [owned, tag, val] = dispatchBuiltin(f, arity); + + for (uint8_t cnt = 0; cnt < arity; ++cnt) { + auto [owned, tag, val] = getFromStack(0); + popStack(); + if (owned) { + value::releaseValue(tag, val); + } + } + + pushStack(owned, tag, val); + + break; + } + case Instruction::jmp: { + auto jumpOffset = value::readFromMemory<int>(pcPointer); + pcPointer += sizeof(jumpOffset); + + pcPointer += jumpOffset; + break; + } + case Instruction::jmpTrue: { + auto jumpOffset = value::readFromMemory<int>(pcPointer); + pcPointer += sizeof(jumpOffset); + + auto [owned, tag, val] = getFromStack(0); + popStack(); + + if (tag == value::TypeTags::Boolean && val) { + pcPointer += jumpOffset; + } + + if (owned) { + value::releaseValue(tag, val); + } + break; + } + case Instruction::jmpNothing: { + auto jumpOffset = value::readFromMemory<int>(pcPointer); + pcPointer += sizeof(jumpOffset); + + auto [owned, tag, val] = getFromStack(0); + if (tag == value::TypeTags::Nothing) { + pcPointer += jumpOffset; + } + break; + } + case Instruction::fail: { + auto [ownedCode, tagCode, valCode] = getFromStack(1); + invariant(tagCode == value::TypeTags::NumberInt64); + + auto [ownedMsg, tagMsg, valMsg] = getFromStack(0); + invariant(value::isString(tagMsg)); + + ErrorCodes::Error code{ + static_cast<ErrorCodes::Error>(value::bitcastTo<int64_t>(valCode))}; + std::string message{value::getStringView(tagMsg, valMsg)}; + + uasserted(code, message); + + break; + } + default: + MONGO_UNREACHABLE; + } + } + } + uassert( + 4822801, "The evaluation stack must hold only a single value", _argStackOwned.size() == 1); + + auto owned = _argStackOwned[0]; + auto tag = _argStackTags[0]; + auto val = _argStackVals[0]; + + _argStackOwned.clear(); + _argStackTags.clear(); + _argStackVals.clear(); + + return {owned, tag, val}; +} + +bool ByteCode::runPredicate(CodeFragment* code) { + auto [owned, tag, val] = run(code); + + bool pass = (tag == value::TypeTags::Boolean) && (val != 0); + + if (owned) { + value::releaseValue(tag, val); + } + + return pass; +} +} // namespace vm +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h new file mode 100644 index 00000000000..63909c856db --- /dev/null +++ b/src/mongo/db/exec/sbe/vm/vm.h @@ -0,0 +1,407 @@ +/** + * Copyright (C) 2019-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <cstdint> +#include <memory> +#include <vector> + +#include "mongo/db/exec/sbe/values/value.h" + +namespace mongo { +namespace sbe { +namespace vm { +template <typename Op> +std::pair<value::TypeTags, value::Value> genericNumericCompare(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue, + Op op) { + + if (value::isNumber(lhsTag) && value::isNumber(rhsTag)) { + switch (getWidestNumericalType(lhsTag, rhsTag)) { + case value::TypeTags::NumberInt32: { + auto result = op(value::numericCast<int32_t>(lhsTag, lhsValue), + value::numericCast<int32_t>(rhsTag, rhsValue)); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberInt64: { + auto result = op(value::numericCast<int64_t>(lhsTag, lhsValue), + value::numericCast<int64_t>(rhsTag, rhsValue)); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDouble: { + auto result = op(value::numericCast<double>(lhsTag, lhsValue), + value::numericCast<double>(rhsTag, rhsValue)); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } + case value::TypeTags::NumberDecimal: { + auto result = op(value::numericCast<Decimal128>(lhsTag, lhsValue), + value::numericCast<Decimal128>(rhsTag, rhsValue)); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } + default: + MONGO_UNREACHABLE; + } + } else if (isString(lhsTag) && isString(rhsTag)) { + auto lhsStr = getStringView(lhsTag, lhsValue); + auto rhsStr = getStringView(rhsTag, rhsValue); + auto result = op(lhsStr.compare(rhsStr), 0); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } else if (lhsTag == value::TypeTags::Date && rhsTag == value::TypeTags::Date) { + auto result = op(value::bitcastTo<int64_t>(lhsValue), value::bitcastTo<int64_t>(rhsValue)); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } else if (lhsTag == value::TypeTags::Timestamp && rhsTag == value::TypeTags::Timestamp) { + auto result = + op(value::bitcastTo<uint64_t>(lhsValue), value::bitcastTo<uint64_t>(rhsValue)); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } else if (lhsTag == value::TypeTags::Boolean && rhsTag == value::TypeTags::Boolean) { + auto result = op(lhsValue != 0, rhsValue != 0); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } else if (lhsTag == value::TypeTags::Null && rhsTag == value::TypeTags::Null) { + // This is where Mongo differs from SQL. + auto result = op(0, 0); + return {value::TypeTags::Boolean, value::bitcastFrom(result)}; + } + + return {value::TypeTags::Nothing, 0}; +} + +struct Instruction { + enum Tags { + pushConstVal, + pushAccessVal, + pushMoveVal, + pushLocalVal, + pop, + swap, + + add, + sub, + mul, + div, + negate, + + logicNot, + + less, + lessEq, + greater, + greaterEq, + eq, + neq, + + // 3 way comparison (spaceship) with bson woCompare semantics. + cmp3w, + + fillEmpty, + + getField, + + aggSum, + aggMin, + aggMax, + aggFirst, + aggLast, + + exists, + isNull, + isObject, + isArray, + isString, + isNumber, + + function, + + jmp, // offset is calculated from the end of instruction + jmpTrue, + jmpNothing, + + fail, + + lastInstruction // this is just a marker used to calculate number of instructions + }; + + // Make sure that values in this arrays are always in-sync with the enum. + static int stackOffset[]; + + uint8_t tag; +}; +static_assert(sizeof(Instruction) == sizeof(uint8_t)); + +enum class Builtin : uint8_t { + split, + regexMatch, + dropFields, + newObj, + ksToString, // KeyString to string + newKs, // new KeyString + abs, // absolute value + addToArray, // agg function to append to an array + addToSet, // agg function to append to a set +}; + +class CodeFragment { +public: + auto& instrs() { + return _instrs; + } + auto stackSize() const { + return _stackSize; + } + void removeFixup(FrameId frameId); + + void append(std::unique_ptr<CodeFragment> code); + void append(std::unique_ptr<CodeFragment> lhs, std::unique_ptr<CodeFragment> rhs); + void appendConstVal(value::TypeTags tag, value::Value val); + void appendAccessVal(value::SlotAccessor* accessor); + void appendMoveVal(value::SlotAccessor* accessor); + void appendLocalVal(FrameId frameId, int stackOffset); + void appendPop() { + appendSimpleInstruction(Instruction::pop); + } + void appendSwap() { + appendSimpleInstruction(Instruction::swap); + } + void appendAdd(); + void appendSub(); + void appendMul(); + void appendDiv(); + void appendNegate(); + void appendNot(); + void appendLess() { + appendSimpleInstruction(Instruction::less); + } + void appendLessEq() { + appendSimpleInstruction(Instruction::lessEq); + } + void appendGreater() { + appendSimpleInstruction(Instruction::greater); + } + void appendGreaterEq() { + appendSimpleInstruction(Instruction::greaterEq); + } + void appendEq() { + appendSimpleInstruction(Instruction::eq); + } + void appendNeq() { + appendSimpleInstruction(Instruction::neq); + } + void appendCmp3w() { + appendSimpleInstruction(Instruction::cmp3w); + } + void appendFillEmpty() { + appendSimpleInstruction(Instruction::fillEmpty); + } + void appendGetField(); + void appendSum(); + void appendMin(); + void appendMax(); + void appendFirst(); + void appendLast(); + void appendExists(); + void appendIsNull(); + void appendIsObject(); + void appendIsArray(); + void appendIsString(); + void appendIsNumber(); + void appendFunction(Builtin f, uint8_t arity); + void appendJump(int jumpOffset); + void appendJumpTrue(int jumpOffset); + void appendJumpNothing(int jumpOffset); + void appendFail() { + appendSimpleInstruction(Instruction::fail); + } + +private: + void appendSimpleInstruction(Instruction::Tags tag); + auto allocateSpace(size_t size) { + auto oldSize = _instrs.size(); + _instrs.resize(oldSize + size); + return _instrs.data() + oldSize; + } + + void adjustStackSimple(const Instruction& i); + void fixup(int offset); + void copyCodeAndFixup(const CodeFragment& from); + + std::vector<uint8_t> _instrs; + + /** + * Local variables bound by the let expressions live on the stack and are accessed by knowing an + * offset from the top of the stack. As CodeFragments are appened together the offsets must be + * fixed up to account for movement of the top of the stack. + * The FixUp structure holds a "pointer" to the bytecode where we have to adjust the stack + * offset. + */ + struct FixUp { + FrameId frameId; + size_t offset; + }; + std::vector<FixUp> _fixUps; + + int _stackSize{0}; +}; + +class ByteCode { +public: + ~ByteCode(); + + std::tuple<uint8_t, value::TypeTags, value::Value> run(CodeFragment* code); + bool runPredicate(CodeFragment* code); + +private: + std::vector<uint8_t> _argStackOwned; + std::vector<value::TypeTags> _argStackTags; + std::vector<value::Value> _argStackVals; + + std::tuple<bool, value::TypeTags, value::Value> genericAdd(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue); + std::tuple<bool, value::TypeTags, value::Value> genericSub(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue); + std::tuple<bool, value::TypeTags, value::Value> genericMul(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue); + std::tuple<bool, value::TypeTags, value::Value> genericDiv(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue); + std::tuple<bool, value::TypeTags, value::Value> genericAbs(value::TypeTags operandTag, + value::Value operandValue); + std::tuple<bool, value::TypeTags, value::Value> genericNot(value::TypeTags tag, + value::Value value); + template <typename Op> + std::pair<value::TypeTags, value::Value> genericCompare(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue, + Op op = {}) { + return genericNumericCompare(lhsTag, lhsValue, rhsTag, rhsValue, op); + } + + std::pair<value::TypeTags, value::Value> genericCompareEq(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue); + + std::pair<value::TypeTags, value::Value> genericCompareNeq(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue); + + std::pair<value::TypeTags, value::Value> compare3way(value::TypeTags lhsTag, + value::Value lhsValue, + value::TypeTags rhsTag, + value::Value rhsValue); + + std::tuple<bool, value::TypeTags, value::Value> getField(value::TypeTags objTag, + value::Value objValue, + value::TypeTags fieldTag, + value::Value fieldValue); + + std::tuple<bool, value::TypeTags, value::Value> aggSum(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue); + + std::tuple<bool, value::TypeTags, value::Value> aggMin(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue); + + std::tuple<bool, value::TypeTags, value::Value> aggMax(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue); + + std::tuple<bool, value::TypeTags, value::Value> aggFirst(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue); + + std::tuple<bool, value::TypeTags, value::Value> aggLast(value::TypeTags accTag, + value::Value accValue, + value::TypeTags fieldTag, + value::Value fieldValue); + + std::tuple<bool, value::TypeTags, value::Value> builtinSplit(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinRegexMatch(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinDropFields(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinNewObj(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinKeyStringToString(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinNewKeyString(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinAbs(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinAddToArray(uint8_t arity); + std::tuple<bool, value::TypeTags, value::Value> builtinAddToSet(uint8_t arity); + + std::tuple<bool, value::TypeTags, value::Value> dispatchBuiltin(Builtin f, uint8_t arity); + + std::tuple<bool, value::TypeTags, value::Value> getFromStack(size_t offset) { + auto backOffset = _argStackOwned.size() - 1 - offset; + auto owned = _argStackOwned[backOffset]; + auto tag = _argStackTags[backOffset]; + auto val = _argStackVals[backOffset]; + + return {owned, tag, val}; + } + + void setStack(size_t offset, bool owned, value::TypeTags tag, value::Value val) { + auto backOffset = _argStackOwned.size() - 1 - offset; + _argStackOwned[backOffset] = owned; + _argStackTags[backOffset] = tag; + _argStackVals[backOffset] = val; + } + + void pushStack(bool owned, value::TypeTags tag, value::Value val) { + _argStackOwned.push_back(owned); + _argStackTags.push_back(tag); + _argStackVals.push_back(val); + } + + void topStack(bool owned, value::TypeTags tag, value::Value val) { + _argStackOwned.back() = owned; + _argStackTags.back() = tag; + _argStackVals.back() = val; + } + + void popStack() { + _argStackOwned.pop_back(); + _argStackTags.pop_back(); + _argStackVals.pop_back(); + } +}; +} // namespace vm +} // namespace sbe +} // namespace mongo diff --git a/src/mongo/db/exec/stagedebug_cmd.cpp b/src/mongo/db/exec/stagedebug_cmd.cpp index 31556bc0217..4442f69a597 100644 --- a/src/mongo/db/exec/stagedebug_cmd.cpp +++ b/src/mongo/db/exec/stagedebug_cmd.cpp @@ -182,8 +182,11 @@ public: unique_ptr<PlanStage> rootFetch = std::make_unique<FetchStage>( expCtx.get(), ws.get(), std::move(userRoot), nullptr, collection); - auto statusWithPlanExecutor = PlanExecutor::make( - expCtx, std::move(ws), std::move(rootFetch), collection, PlanExecutor::YIELD_AUTO); + auto statusWithPlanExecutor = PlanExecutor::make(expCtx, + std::move(ws), + std::move(rootFetch), + collection, + PlanYieldPolicy::YieldPolicy::YIELD_AUTO); fassert(28536, statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); diff --git a/src/mongo/db/exec/subplan.cpp b/src/mongo/db/exec/subplan.cpp index 7735e7136b4..b8abd0d620c 100644 --- a/src/mongo/db/exec/subplan.cpp +++ b/src/mongo/db/exec/subplan.cpp @@ -37,6 +37,7 @@ #include <vector> #include "mongo/db/exec/multi_plan.h" +#include "mongo/db/exec/plan_cache_util.h" #include "mongo/db/exec/scoped_timer.h" #include "mongo/db/matcher/extensions_callback_real.h" #include "mongo/db/query/collection_query_info.h" @@ -44,9 +45,8 @@ #include "mongo/db/query/plan_executor.h" #include "mongo/db/query/planner_access.h" #include "mongo/db/query/planner_analysis.h" -#include "mongo/db/query/query_planner.h" #include "mongo/db/query/query_planner_common.h" -#include "mongo/db/query/stage_builder.h" +#include "mongo/db/query/stage_builder_util.h" #include "mongo/logv2/log.h" #include "mongo/util/scopeguard.h" #include "mongo/util/transitional_tools_do_not_use/vector_spooling.h" @@ -104,271 +104,6 @@ bool SubplanStage::canUseSubplanning(const CanonicalQuery& query) { return MatchExpression::OR == expr->matchType() && expr->numChildren() > 0; } -Status SubplanStage::planSubqueries() { - _orExpression = _query->root()->shallowClone(); - for (size_t i = 0; i < _plannerParams.indices.size(); ++i) { - const IndexEntry& ie = _plannerParams.indices[i]; - const auto insertionRes = _indexMap.insert(std::make_pair(ie.identifier, i)); - // Be sure the key was not already in the map. - invariant(insertionRes.second); - LOGV2_DEBUG(20598, 5, "Subplanner: index {i} is {ie}", "i"_attr = i, "ie"_attr = ie); - } - - for (size_t i = 0; i < _orExpression->numChildren(); ++i) { - // We need a place to shove the results from planning this branch. - _branchResults.push_back(std::make_unique<BranchPlanningResult>()); - BranchPlanningResult* branchResult = _branchResults.back().get(); - - MatchExpression* orChild = _orExpression->getChild(i); - - // Turn the i-th child into its own query. - auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), *_query, orChild); - if (!statusWithCQ.isOK()) { - str::stream ss; - ss << "Can't canonicalize subchild " << orChild->debugString() << " " - << statusWithCQ.getStatus().reason(); - return Status(ErrorCodes::BadValue, ss); - } - - branchResult->canonicalQuery = std::move(statusWithCQ.getValue()); - - // Plan the i-th child. We might be able to find a plan for the i-th child in the plan - // cache. If there's no cached plan, then we generate and rank plans using the MPS. - const auto* planCache = CollectionQueryInfo::get(collection()).getPlanCache(); - - // Populate branchResult->cachedSolution if an active cachedSolution entry exists. - if (planCache->shouldCacheQuery(*branchResult->canonicalQuery)) { - auto planCacheKey = planCache->computeKey(*branchResult->canonicalQuery); - if (auto cachedSol = planCache->getCacheEntryIfActive(planCacheKey)) { - // We have a CachedSolution. Store it for later. - LOGV2_DEBUG( - 20599, - 5, - "Subplanner: cached plan found for child {i} of {orExpression_numChildren}", - "i"_attr = i, - "orExpression_numChildren"_attr = _orExpression->numChildren()); - - branchResult->cachedSolution = std::move(cachedSol); - } - } - - if (!branchResult->cachedSolution) { - // No CachedSolution found. We'll have to plan from scratch. - LOGV2_DEBUG(20600, - 5, - "Subplanner: planning child {i} of {orExpression_numChildren}", - "i"_attr = i, - "orExpression_numChildren"_attr = _orExpression->numChildren()); - - // We don't set NO_TABLE_SCAN because peeking at the cache data will keep us from - // considering any plan that's a collscan. - invariant(branchResult->solutions.empty()); - auto solutions = QueryPlanner::plan(*branchResult->canonicalQuery, _plannerParams); - if (!solutions.isOK()) { - str::stream ss; - ss << "Can't plan for subchild " << branchResult->canonicalQuery->toString() << " " - << solutions.getStatus().reason(); - return Status(ErrorCodes::BadValue, ss); - } - branchResult->solutions = std::move(solutions.getValue()); - - LOGV2_DEBUG(20601, - 5, - "Subplanner: got {branchResult_solutions_size} solutions", - "branchResult_solutions_size"_attr = branchResult->solutions.size()); - } - } - - return Status::OK(); -} - -namespace { - -/** - * On success, applies the index tags from 'branchCacheData' (which represent the winning - * plan for 'orChild') to 'compositeCacheData'. - */ -Status tagOrChildAccordingToCache(PlanCacheIndexTree* compositeCacheData, - SolutionCacheData* branchCacheData, - MatchExpression* orChild, - const std::map<IndexEntry::Identifier, size_t>& indexMap) { - invariant(compositeCacheData); - - // We want a well-formed *indexed* solution. - if (nullptr == branchCacheData) { - // For example, we don't cache things for 2d indices. - str::stream ss; - ss << "No cache data for subchild " << orChild->debugString(); - return Status(ErrorCodes::NoQueryExecutionPlans, ss); - } - - if (SolutionCacheData::USE_INDEX_TAGS_SOLN != branchCacheData->solnType) { - str::stream ss; - ss << "No indexed cache data for subchild " << orChild->debugString(); - return Status(ErrorCodes::NoQueryExecutionPlans, ss); - } - - // Add the index assignments to our original query. - Status tagStatus = - QueryPlanner::tagAccordingToCache(orChild, branchCacheData->tree.get(), indexMap); - - if (!tagStatus.isOK()) { - str::stream ss; - ss << "Failed to extract indices from subchild " << orChild->debugString(); - return tagStatus.withContext(ss); - } - - // Add the child's cache data to the cache data we're creating for the main query. - compositeCacheData->children.push_back(branchCacheData->tree->clone()); - - return Status::OK(); -} - -} // namespace - -Status SubplanStage::choosePlanForSubqueries(PlanYieldPolicy* yieldPolicy) { - // This is the skeleton of index selections that is inserted into the cache. - std::unique_ptr<PlanCacheIndexTree> cacheData(new PlanCacheIndexTree()); - - for (size_t i = 0; i < _orExpression->numChildren(); ++i) { - MatchExpression* orChild = _orExpression->getChild(i); - BranchPlanningResult* branchResult = _branchResults[i].get(); - - if (branchResult->cachedSolution.get()) { - // We can get the index tags we need out of the cache. - Status tagStatus = tagOrChildAccordingToCache( - cacheData.get(), branchResult->cachedSolution->plannerData[0], orChild, _indexMap); - if (!tagStatus.isOK()) { - return tagStatus; - } - } else if (1 == branchResult->solutions.size()) { - QuerySolution* soln = branchResult->solutions.front().get(); - Status tagStatus = tagOrChildAccordingToCache( - cacheData.get(), soln->cacheData.get(), orChild, _indexMap); - if (!tagStatus.isOK()) { - return tagStatus; - } - } else { - // N solutions, rank them. - - // We already checked for zero solutions in planSubqueries(...). - invariant(!branchResult->solutions.empty()); - - _ws->clear(); - - // We pass the SometimesCache option to the MPS because the SubplanStage currently does - // not use the CachedPlanStage's eviction mechanism. We therefore are more conservative - // about putting a potentially bad plan into the cache in the subplan path. We - // temporarily add the MPS to _children to ensure that we pass down all save/restore - // messages that can be generated if pickBestPlan yields. - invariant(_children.empty()); - _children.emplace_back( - std::make_unique<MultiPlanStage>(expCtx(), - collection(), - branchResult->canonicalQuery.get(), - MultiPlanStage::CachingMode::SometimesCache)); - ON_BLOCK_EXIT([&] { - invariant(_children.size() == 1); // Make sure nothing else was added to _children. - _children.pop_back(); - }); - MultiPlanStage* multiPlanStage = static_cast<MultiPlanStage*>(child().get()); - - // Dump all the solutions into the MPS. - for (size_t ix = 0; ix < branchResult->solutions.size(); ++ix) { - auto nextPlanRoot = StageBuilder::build(opCtx(), - collection(), - *branchResult->canonicalQuery, - *branchResult->solutions[ix], - _ws); - - multiPlanStage->addPlan( - std::move(branchResult->solutions[ix]), std::move(nextPlanRoot), _ws); - } - - Status planSelectStat = multiPlanStage->pickBestPlan(yieldPolicy); - if (!planSelectStat.isOK()) { - return planSelectStat; - } - - if (!multiPlanStage->bestPlanChosen()) { - str::stream ss; - ss << "Failed to pick best plan for subchild " - << branchResult->canonicalQuery->toString(); - return Status(ErrorCodes::NoQueryExecutionPlans, ss); - } - - QuerySolution* bestSoln = multiPlanStage->bestSolution(); - - // Check that we have good cache data. For example, we don't cache things - // for 2d indices. - if (nullptr == bestSoln->cacheData.get()) { - str::stream ss; - ss << "No cache data for subchild " << orChild->debugString(); - return Status(ErrorCodes::NoQueryExecutionPlans, ss); - } - - if (SolutionCacheData::USE_INDEX_TAGS_SOLN != bestSoln->cacheData->solnType) { - str::stream ss; - ss << "No indexed cache data for subchild " << orChild->debugString(); - return Status(ErrorCodes::NoQueryExecutionPlans, ss); - } - - // Add the index assignments to our original query. - Status tagStatus = QueryPlanner::tagAccordingToCache( - orChild, bestSoln->cacheData->tree.get(), _indexMap); - - if (!tagStatus.isOK()) { - str::stream ss; - ss << "Failed to extract indices from subchild " << orChild->debugString(); - return tagStatus.withContext(ss); - } - - cacheData->children.push_back(bestSoln->cacheData->tree->clone()); - } - } - - // Must do this before using the planner functionality. - prepareForAccessPlanning(_orExpression.get()); - - // Use the cached index assignments to build solnRoot. Takes ownership of '_orExpression'. - std::unique_ptr<QuerySolutionNode> solnRoot(QueryPlannerAccess::buildIndexedDataAccess( - *_query, std::move(_orExpression), _plannerParams.indices, _plannerParams)); - - if (!solnRoot) { - str::stream ss; - ss << "Failed to build indexed data path for subplanned query\n"; - return Status(ErrorCodes::NoQueryExecutionPlans, ss); - } - - LOGV2_DEBUG(20602, - 5, - "Subplanner: fully tagged tree is {solnRoot}", - "solnRoot"_attr = redact(solnRoot->toString())); - - _compositeSolution = - QueryPlannerAnalysis::analyzeDataAccess(*_query, _plannerParams, std::move(solnRoot)); - - if (nullptr == _compositeSolution.get()) { - str::stream ss; - ss << "Failed to analyze subplanned query"; - return Status(ErrorCodes::NoQueryExecutionPlans, ss); - } - - LOGV2_DEBUG(20603, - 5, - "Subplanner: Composite solution is {compositeSolution}", - "compositeSolution"_attr = redact(_compositeSolution->toString())); - - // Use the index tags from planning each branch to construct the composite solution, - // and set that solution as our child stage. - _ws->clear(); - auto root = StageBuilder::build(opCtx(), collection(), *_query, *_compositeSolution.get(), _ws); - invariant(_children.empty()); - _children.emplace_back(std::move(root)); - - return Status::OK(); -} - Status SubplanStage::choosePlanWholeQuery(PlanYieldPolicy* yieldPolicy) { // Clear out the working set. We'll start with a fresh working set. _ws->clear(); @@ -384,7 +119,8 @@ Status SubplanStage::choosePlanWholeQuery(PlanYieldPolicy* yieldPolicy) { if (1 == solutions.size()) { // Only one possible plan. Run it. Build the stages from the solution. - auto root = StageBuilder::build(opCtx(), collection(), *_query, *solutions[0], _ws); + auto&& root = stage_builder::buildClassicExecutableTree( + expCtx()->opCtx, collection(), *_query, *solutions[0], _ws); invariant(_children.empty()); _children.emplace_back(std::move(root)); @@ -405,9 +141,8 @@ Status SubplanStage::choosePlanWholeQuery(PlanYieldPolicy* yieldPolicy) { solutions[ix]->cacheData->indexFilterApplied = _plannerParams.indexFiltersApplied; } - auto nextPlanRoot = - StageBuilder::build(opCtx(), collection(), *_query, *solutions[ix], _ws); - + auto&& nextPlanRoot = stage_builder::buildClassicExecutableTree( + expCtx()->opCtx, collection(), *_query, *solutions[ix], _ws); multiPlanStage->addPlan(std::move(solutions[ix]), std::move(nextPlanRoot), _ws); } @@ -435,24 +170,85 @@ Status SubplanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { ON_BLOCK_EXIT([this] { releaseAllIndicesRequirement(); }); // Plan each branch of the $or. - Status subplanningStatus = planSubqueries(); + auto subplanningStatus = + QueryPlanner::planSubqueries(expCtx()->opCtx, + collection(), + CollectionQueryInfo::get(collection()).getPlanCache(), + *_query, + _plannerParams); if (!subplanningStatus.isOK()) { return choosePlanWholeQuery(yieldPolicy); } + // Remember whether each branch of the $or was planned from a cached solution. + auto subplanningResult = std::move(subplanningStatus.getValue()); + _branchPlannedFromCache.clear(); + for (auto&& branch : subplanningResult.branches) { + _branchPlannedFromCache.push_back(branch->cachedSolution != nullptr); + } + // Use the multi plan stage to select a winning plan for each branch, and then construct // the overall winning plan from the resulting index tags. - Status subplanSelectStat = choosePlanForSubqueries(yieldPolicy); + auto multiplanCallback = [&](CanonicalQuery* cq, + std::vector<std::unique_ptr<QuerySolution>> solutions) + -> StatusWith<std::unique_ptr<QuerySolution>> { + _ws->clear(); + + // We pass the SometimesCache option to the MPS because the SubplanStage currently does + // not use the CachedPlanStage's eviction mechanism. We therefore are more conservative + // about putting a potentially bad plan into the cache in the subplan path. + // + // We temporarily add the MPS to _children to ensure that we pass down all save/restore + // messages that can be generated if pickBestPlan yields. + invariant(_children.empty()); + _children.emplace_back(std::make_unique<MultiPlanStage>( + expCtx(), collection(), cq, PlanCachingMode::SometimesCache)); + ON_BLOCK_EXIT([&] { + invariant(_children.size() == 1); // Make sure nothing else was added to _children. + _children.pop_back(); + }); + MultiPlanStage* multiPlanStage = static_cast<MultiPlanStage*>(child().get()); + + // Dump all the solutions into the MPS. + for (size_t ix = 0; ix < solutions.size(); ++ix) { + auto&& nextPlanRoot = stage_builder::buildClassicExecutableTree( + expCtx()->opCtx, collection(), *cq, *solutions[ix], _ws); + + multiPlanStage->addPlan(std::move(solutions[ix]), std::move(nextPlanRoot), _ws); + } + + Status planSelectStat = multiPlanStage->pickBestPlan(yieldPolicy); + if (!planSelectStat.isOK()) { + return planSelectStat; + } + + if (!multiPlanStage->bestPlanChosen()) { + str::stream ss; + ss << "Failed to pick best plan for subchild " << cq->toString(); + return Status(ErrorCodes::NoQueryExecutionPlans, ss); + } + return multiPlanStage->bestSolution(); + }; + auto subplanSelectStat = QueryPlanner::choosePlanForSubqueries( + *_query, _plannerParams, std::move(subplanningResult), multiplanCallback); if (!subplanSelectStat.isOK()) { if (subplanSelectStat != ErrorCodes::NoQueryExecutionPlans) { // Query planning can continue if we failed to find a solution for one of the // children. Otherwise, it cannot, as it may no longer be safe to access the collection // (and index may have been dropped, we may have exceeded the time limit, etc). - return subplanSelectStat; + return subplanSelectStat.getStatus(); } return choosePlanWholeQuery(yieldPolicy); } + // Build a plan stage tree from the the composite solution and add it as our child stage. + _compositeSolution = std::move(subplanSelectStat.getValue()); + invariant(_children.empty()); + auto&& root = stage_builder::buildClassicExecutableTree( + expCtx()->opCtx, collection(), *_query, *_compositeSolution, _ws); + _children.emplace_back(std::move(root)); + _ws->clear(); + return Status::OK(); } @@ -478,12 +274,7 @@ unique_ptr<PlanStageStats> SubplanStage::getStats() { return ret; } -bool SubplanStage::branchPlannedFromCache(size_t i) const { - return nullptr != _branchResults[i]->cachedSolution.get(); -} - const SpecificStats* SubplanStage::getSpecificStats() const { return nullptr; } - } // namespace mongo diff --git a/src/mongo/db/exec/subplan.h b/src/mongo/db/exec/subplan.h index 07d8f956ca8..6b5d5b7448b 100644 --- a/src/mongo/db/exec/subplan.h +++ b/src/mongo/db/exec/subplan.h @@ -40,6 +40,7 @@ #include "mongo/db/query/canonical_query.h" #include "mongo/db/query/plan_cache.h" #include "mongo/db/query/plan_yield_policy.h" +#include "mongo/db/query/query_planner.h" #include "mongo/db/query/query_planner_params.h" #include "mongo/db/query/query_solution.h" #include "mongo/db/record_id.h" @@ -115,7 +116,10 @@ public: * Returns true if the i-th branch was planned by retrieving a cached solution, * otherwise returns false. */ - bool branchPlannedFromCache(size_t i) const; + bool branchPlannedFromCache(size_t i) const { + invariant(i < _branchPlannedFromCache.size()); + return _branchPlannedFromCache[i]; + } /** * Provide access to the query solution for our composite solution. Does not relinquish @@ -127,48 +131,6 @@ public: private: /** - * A class used internally in order to keep track of the results of planning - * a particular $or branch. - */ - struct BranchPlanningResult { - BranchPlanningResult(const BranchPlanningResult&) = delete; - BranchPlanningResult& operator=(const BranchPlanningResult&) = delete; - - public: - BranchPlanningResult() {} - - // A parsed version of one branch of the $or. - std::unique_ptr<CanonicalQuery> canonicalQuery; - - // If there is cache data available, then we store it here rather than generating - // a set of alternate plans for the branch. The index tags from the cache data - // can be applied directly to the parent $or MatchExpression when generating the - // composite solution. - std::unique_ptr<CachedSolution> cachedSolution; - - // Query solutions resulting from planning the $or branch. - std::vector<std::unique_ptr<QuerySolution>> solutions; - }; - - /** - * Plan each branch of the $or independently, and store the resulting - * lists of query solutions in '_solutions'. - * - * Called from SubplanStage::make so that construction of the subplan stage - * fails immediately, rather than returning a plan executor and subsequently - * through getNext(...). - */ - Status planSubqueries(); - - /** - * Uses the query planning results from planSubqueries() and the multi plan stage - * to select the best plan for each branch. - * - * Helper for pickBestPlan(). - */ - Status choosePlanForSubqueries(PlanYieldPolicy* yieldPolicy); - - /** * Used as a fallback if subplanning fails. Helper for pickBestPlan(). */ Status choosePlanWholeQuery(PlanYieldPolicy* yieldPolicy); @@ -181,20 +143,11 @@ private: // Not owned here. CanonicalQuery* _query; - // The copy of the query that we will annotate with tags and use to construct the composite - // solution. Must be a rooted $or query, or a contained $or that has been rewritten to a - // rooted $or. - std::unique_ptr<MatchExpression> _orExpression; - // If we successfully create a "composite solution" by planning each $or branch // independently, that solution is owned here. std::unique_ptr<QuerySolution> _compositeSolution; - // Holds a list of the results from planning each branch. - std::vector<std::unique_ptr<BranchPlanningResult>> _branchResults; - - // We need this to extract cache-friendly index data from the index assignments. - std::map<IndexEntry::Identifier, size_t> _indexMap; + // Indicates whether i-th branch of the rooted $or query was planned from a cached solution. + std::vector<bool> _branchPlannedFromCache; }; - } // namespace mongo diff --git a/src/mongo/db/exec/trial_period_utils.cpp b/src/mongo/db/exec/trial_period_utils.cpp new file mode 100644 index 00000000000..759e41b3577 --- /dev/null +++ b/src/mongo/db/exec/trial_period_utils.cpp @@ -0,0 +1,66 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/trial_period_utils.h" + +#include "mongo/db/catalog/collection.h" + +namespace mongo::trial_period { +size_t getTrialPeriodMaxWorks(OperationContext* opCtx, const Collection* collection) { + // Run each plan some number of times. This number is at least as great as + // 'internalQueryPlanEvaluationWorks', but may be larger for big collections. + size_t numWorks = internalQueryPlanEvaluationWorks.load(); + if (nullptr != collection) { + // For large collections, the number of works is set to be this fraction of the collection + // size. + double fraction = internalQueryPlanEvaluationCollFraction; + + numWorks = std::max(static_cast<size_t>(internalQueryPlanEvaluationWorks.load()), + static_cast<size_t>(fraction * collection->numRecords(opCtx))); + } + + return numWorks; +} + +size_t getTrialPeriodNumToReturn(const CanonicalQuery& query) { + // Determine the number of results which we will produce during the plan ranking phase before + // stopping. + size_t numResults = static_cast<size_t>(internalQueryPlanEvaluationMaxResults.load()); + if (query.getQueryRequest().getNToReturn()) { + numResults = + std::min(static_cast<size_t>(*query.getQueryRequest().getNToReturn()), numResults); + } else if (query.getQueryRequest().getLimit()) { + numResults = std::min(static_cast<size_t>(*query.getQueryRequest().getLimit()), numResults); + } + + return numResults; +} +} // namespace mongo::trial_period diff --git a/src/mongo/db/exec/trial_period_utils.h b/src/mongo/db/exec/trial_period_utils.h new file mode 100644 index 00000000000..4919e8be1c9 --- /dev/null +++ b/src/mongo/db/exec/trial_period_utils.h @@ -0,0 +1,51 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include "mongo/db/query/canonical_query.h" + +namespace mongo { +class Collection; + +namespace trial_period { +/** + * Returns the number of times that we are willing to work a plan during a trial period. + * + * Calculated based on a fixed query knob and the size of the collection. + */ +size_t getTrialPeriodMaxWorks(OperationContext* opCtx, const Collection* collection); + +/** + * Returns the max number of documents which we should allow any plan to return during the + * trial period. As soon as any plan hits this number of documents, the trial period ends. + */ +size_t getTrialPeriodNumToReturn(const CanonicalQuery& query); +} // namespace trial_period +} // namespace mongo diff --git a/src/mongo/db/exec/trial_run_progress_tracker.h b/src/mongo/db/exec/trial_run_progress_tracker.h new file mode 100644 index 00000000000..a9bf9110e24 --- /dev/null +++ b/src/mongo/db/exec/trial_run_progress_tracker.h @@ -0,0 +1,90 @@ +/** + * Copyright (C) 2020-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <cstddef> +#include <cstdint> +#include <type_traits> + +namespace mongo { +/** + * During the runtime planning phase this tracker is used to track the progress of the work done + * so far. A plan stage in a candidate execution tree may supply the number of documents it has + * processed, or the number of physical reads performed, and the tracker will use it to check if + * the execution plan has progressed enough. + */ +class TrialRunProgressTracker final { +public: + /** + * The type of metric which can be collected and tracked during a trial run. + */ + enum TrialRunMetric : uint8_t { + // Number of documents returned during a trial run. + kNumResults, + // Number of physical reads performed during a trial run. Once a storage cursor advances, + // it counts as a single physical read. + kNumReads, + // Must always be the last element to hold the number of element in the enum. + kLastElem + }; + + template <typename... MaxMetrics, + std::enable_if_t<sizeof...(MaxMetrics) == TrialRunMetric::kLastElem, int> = 0> + TrialRunProgressTracker(MaxMetrics... maxMetrics) : _maxMetrics{maxMetrics...} {} + + /** + * Increments the trial run metric specified as a template parameter 'metric' by the + * 'metricIncrement' value and returns 'true' if the updated metric value has reached + * its maximum. + * + * If the metric has already reached its maximum value before this call, this method + * returns 'true' immediately without incrementing the metric. + */ + template <TrialRunMetric metric> + bool trackProgress(size_t metricIncrement) { + static_assert(metric >= 0 && metric < sizeof(_metrics)); + + if (_done) { + return true; + } + + _metrics[metric] += metricIncrement; + if (_metrics[metric] >= _maxMetrics[metric]) { + _done = true; + } + return _done; + } + +private: + const size_t _maxMetrics[TrialRunMetric::kLastElem]; + size_t _metrics[TrialRunMetric::kLastElem]{0}; + bool _done{false}; +}; +} // namespace mongo diff --git a/src/mongo/db/exec/trial_stage.cpp b/src/mongo/db/exec/trial_stage.cpp index 0b7ef73d4bb..0d21af2ac6b 100644 --- a/src/mongo/db/exec/trial_stage.cpp +++ b/src/mongo/db/exec/trial_stage.cpp @@ -76,11 +76,11 @@ Status TrialStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) { while (!_specificStats.trialCompleted) { WorkingSetID id = WorkingSet::INVALID_ID; const bool mustYield = (work(&id) == PlanStage::NEED_YIELD); - if (mustYield || yieldPolicy->shouldYieldOrInterrupt()) { + if (mustYield || yieldPolicy->shouldYieldOrInterrupt(expCtx()->opCtx)) { if (mustYield && !yieldPolicy->canAutoYield()) { throw WriteConflictException(); } - auto yieldStatus = yieldPolicy->yieldOrInterrupt(); + auto yieldStatus = yieldPolicy->yieldOrInterrupt(expCtx()->opCtx); if (!yieldStatus.isOK()) { return yieldStatus; } |