diff options
| author | Martin Neupauer <martin.neupauer@mongodb.com> | 2020-06-11 08:07:39 +0100 |
|---|---|---|
| committer | Evergreen Agent <no-reply@evergreen.mongodb.com> | 2020-06-11 11:17:49 +0000 |
| commit | e3948d4d8817579b6b03618e64e1b9e8cc2ef086 (patch) | |
| tree | 649bef264a16807b269f7b645a8d2312c4442455 | |
| parent | 0af9c85d7e2ba60f592f2d7a9a35217e254e59fb (diff) | |
| download | mongo-e3948d4d8817579b6b03618e64e1b9e8cc2ef086.tar.gz | |
SERVER-48228 Move slot-based execution engine and supporting changes into the master branch
This is an initial commit for the slot-based execution engine (SBE) which contains:
* Implementation of the core slot-based engine.
* The SBE stage builder, which is responsible for translating a QuerySolution tree
into an SBE plan.
* Other changes necessary for integration with the find command.
Co-authored-by: Anton Korshunov <anton.korshunov@mongodb.com>
Co-authored-by: Justin Seyster <justin.seyster@mongodb.com>
Co-authored-by: David Storch <david.storch@mongodb.com>
242 files changed, 23314 insertions, 2112 deletions
diff --git a/src/mongo/base/error_codes.yml b/src/mongo/base/error_codes.yml index 93f4842e92a..cc97ab45b5d 100644 --- a/src/mongo/base/error_codes.yml +++ b/src/mongo/base/error_codes.yml @@ -370,6 +370,9 @@ error_codes: # The code below is for internal use only and must never be returned in a network response - {code: 315,name: TooStaleToSyncFromSource} + # The code below is for internal use only and must never be returned in a network response. + - {code: 316,name: QueryTrialRunCompleted} + # Error codes 4000-8999 are reserved. # Non-sequential error codes for compatibility only) diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript index 1a7b398efeb..be72fd9ff19 100644 --- a/src/mongo/db/SConscript +++ b/src/mongo/db/SConscript @@ -1080,6 +1080,7 @@ env.Library( 'exec/near.cpp', 'exec/or.cpp', 'exec/pipeline_proxy.cpp', + 'exec/plan_cache_util.cpp', 'exec/plan_stage.cpp', 'exec/projection.cpp', 'exec/queued_data_stage.cpp', @@ -1097,6 +1098,7 @@ env.Library( 'exec/text.cpp', 'exec/text_match.cpp', 'exec/text_or.cpp', + 'exec/trial_period_utils.cpp', 'exec/trial_stage.cpp', 'exec/update_stage.cpp', 'exec/upsert_stage.cpp', @@ -1107,14 +1109,28 @@ env.Library( 'pipeline/document_source_cursor.cpp', 'pipeline/document_source_geo_near_cursor.cpp', 'pipeline/pipeline_d.cpp', + 'query/classic_stage_builder.cpp', 'query/explain.cpp', 'query/find.cpp', 'query/get_executor.cpp', 'query/internal_plans.cpp', 'query/plan_executor_impl.cpp', + 'query/plan_executor_sbe.cpp', 'query/plan_ranker.cpp', - 'query/plan_yield_policy.cpp', - 'query/stage_builder.cpp', + 'query/plan_yield_policy_impl.cpp', + 'query/plan_yield_policy_sbe.cpp', + 'query/sbe_cached_solution_planner.cpp', + 'query/sbe_multi_planner.cpp', + 'query/sbe_plan_ranker.cpp', + 'query/sbe_runtime_planner.cpp', + 'query/sbe_stage_builder.cpp', + 'query/sbe_stage_builder_coll_scan.cpp', + 'query/sbe_stage_builder_expression.cpp', + 'query/sbe_stage_builder_filter.cpp', + 'query/sbe_stage_builder_index_scan.cpp', + 'query/sbe_stage_builder_projection.cpp', + 'query/sbe_sub_planner.cpp', + 'query/stage_builder_util.cpp', 'run_op_kill_cursors.cpp', ], LIBDEPS=[ @@ -1139,6 +1155,7 @@ env.Library( 'db_raii', 'dbdirectclient', 'exec/projection_executor', + 'exec/sbe/query_sbe_storage', 'exec/scoped_timer', 'exec/sort_executor', 'exec/working_set', @@ -1149,6 +1166,7 @@ env.Library( 'matcher/expressions_mongod_only', 'ops/parsed_update', 'pipeline/pipeline', + 'query/plan_yield_policy', 'query/query_common', 'query/query_planner', 'repl/repl_coordinator_interface', diff --git a/src/mongo/db/catalog/capped_utils.cpp b/src/mongo/db/catalog/capped_utils.cpp index 182c7eaa7d0..9efabad72bc 100644 --- a/src/mongo/db/catalog/capped_utils.cpp +++ b/src/mongo/db/catalog/capped_utils.cpp @@ -166,11 +166,12 @@ void cloneCollectionAsCapped(OperationContext* opCtx, long long excessSize = fromCollection->dataSize(opCtx) - allocatedSpaceGuess; - auto exec = InternalPlanner::collectionScan(opCtx, - fromNss.ns(), - fromCollection, - PlanExecutor::WRITE_CONFLICT_RETRY_ONLY, - InternalPlanner::FORWARD); + auto exec = + InternalPlanner::collectionScan(opCtx, + fromNss.ns(), + fromCollection, + PlanYieldPolicy::YieldPolicy::WRITE_CONFLICT_RETRY_ONLY, + InternalPlanner::FORWARD); Snapshotted<BSONObj> objToClone; RecordId loc; diff --git a/src/mongo/db/catalog/collection.h b/src/mongo/db/catalog/collection.h index 920f9ae03a5..c05226fb4ce 100644 --- a/src/mongo/db/catalog/collection.h +++ b/src/mongo/db/catalog/collection.h @@ -543,7 +543,7 @@ public: */ virtual std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> makePlanExecutor( OperationContext* opCtx, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, ScanDirection scanDirection) = 0; virtual void indexBuildSuccess(OperationContext* opCtx, IndexCatalogEntry* index) = 0; diff --git a/src/mongo/db/catalog/collection_impl.cpp b/src/mongo/db/catalog/collection_impl.cpp index 7317c4d9c89..52d8330fa8a 100644 --- a/src/mongo/db/catalog/collection_impl.cpp +++ b/src/mongo/db/catalog/collection_impl.cpp @@ -1190,7 +1190,9 @@ StatusWith<std::vector<BSONObj>> CollectionImpl::addCollationDefaultsToIndexSpec } std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> CollectionImpl::makePlanExecutor( - OperationContext* opCtx, PlanExecutor::YieldPolicy yieldPolicy, ScanDirection scanDirection) { + OperationContext* opCtx, + PlanYieldPolicy::YieldPolicy yieldPolicy, + ScanDirection scanDirection) { auto isForward = scanDirection == ScanDirection::kForward; auto direction = isForward ? InternalPlanner::FORWARD : InternalPlanner::BACKWARD; return InternalPlanner::collectionScan(opCtx, _ns.ns(), this, yieldPolicy, direction); diff --git a/src/mongo/db/catalog/collection_impl.h b/src/mongo/db/catalog/collection_impl.h index 07f6e558d66..de264eea471 100644 --- a/src/mongo/db/catalog/collection_impl.h +++ b/src/mongo/db/catalog/collection_impl.h @@ -351,7 +351,7 @@ public: std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> makePlanExecutor( OperationContext* opCtx, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, ScanDirection scanDirection) final; void indexBuildSuccess(OperationContext* opCtx, IndexCatalogEntry* index) final; diff --git a/src/mongo/db/catalog/collection_mock.h b/src/mongo/db/catalog/collection_mock.h index 547f4555e61..21332edb9e8 100644 --- a/src/mongo/db/catalog/collection_mock.h +++ b/src/mongo/db/catalog/collection_mock.h @@ -270,7 +270,7 @@ public: std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> makePlanExecutor( OperationContext* opCtx, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, ScanDirection scanDirection) { std::abort(); } diff --git a/src/mongo/db/catalog/multi_index_block.cpp b/src/mongo/db/catalog/multi_index_block.cpp index 3b280b728b0..58a696e9f8f 100644 --- a/src/mongo/db/catalog/multi_index_block.cpp +++ b/src/mongo/db/catalog/multi_index_block.cpp @@ -390,11 +390,11 @@ Status MultiIndexBlock::insertAllDocumentsInCollection(OperationContext* opCtx, unsigned long long n = 0; - PlanExecutor::YieldPolicy yieldPolicy; + PlanYieldPolicy::YieldPolicy yieldPolicy; if (isBackgroundBuilding()) { - yieldPolicy = PlanExecutor::YIELD_AUTO; + yieldPolicy = PlanYieldPolicy::YieldPolicy::YIELD_AUTO; } else { - yieldPolicy = PlanExecutor::WRITE_CONFLICT_RETRY_ONLY; + yieldPolicy = PlanYieldPolicy::YieldPolicy::WRITE_CONFLICT_RETRY_ONLY; } auto exec = collection->makePlanExecutor(opCtx, yieldPolicy, Collection::ScanDirection::kForward); diff --git a/src/mongo/db/clientcursor.cpp b/src/mongo/db/clientcursor.cpp index 81da894e302..f7264afade1 100644 --- a/src/mongo/db/clientcursor.cpp +++ b/src/mongo/db/clientcursor.cpp @@ -88,7 +88,6 @@ ClientCursor::ClientCursor(ClientCursorParams params, _originatingCommand(params.originatingCommandObj), _originatingPrivileges(std::move(params.originatingPrivileges)), _queryOptions(params.queryOptions), - _lockPolicy(params.lockPolicy), _needsMerge(params.needsMerge), _exec(std::move(params.exec)), _operationUsingCursor(operationUsingCursor), diff --git a/src/mongo/db/clientcursor.h b/src/mongo/db/clientcursor.h index 5979ed4521c..a9665b878c2 100644 --- a/src/mongo/db/clientcursor.h +++ b/src/mongo/db/clientcursor.h @@ -55,33 +55,12 @@ class RecoveryUnit; * using a CursorManager. See cursor_manager.h for more details. */ struct ClientCursorParams { - // Describes whether callers should acquire locks when using a ClientCursor. Not all cursors - // have the same locking behavior. In particular, find cursors require the caller to lock the - // collection in MODE_IS before calling methods on the underlying plan executor. Aggregate - // cursors, on the other hand, may access multiple collections and acquire their own locks on - // any involved collections while producing query results. Therefore, the caller need not - // explicitly acquire any locks when using a ClientCursor which houses execution machinery for - // an aggregate. - // - // The policy is consulted on getMore in order to determine locking behavior, since during - // getMore we otherwise could not easily know what flavor of cursor we're using. - enum class LockPolicy { - // The caller is responsible for locking the collection over which this ClientCursor - // executes. - kLockExternally, - - // The caller need not hold no locks; this ClientCursor's plan executor acquires any - // necessary locks itself. - kLocksInternally, - }; - ClientCursorParams(std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> planExecutor, NamespaceString nss, UserNameIterator authenticatedUsersIter, WriteConcernOptions writeConcernOptions, repl::ReadConcernArgs readConcernArgs, BSONObj originatingCommandObj, - LockPolicy lockPolicy, PrivilegeVector originatingPrivileges, bool needsMerge) : exec(std::move(planExecutor)), @@ -92,7 +71,6 @@ struct ClientCursorParams { ? exec->getCanonicalQuery()->getQueryRequest().getOptions() : 0), originatingCommandObj(originatingCommandObj.getOwned()), - lockPolicy(lockPolicy), originatingPrivileges(std::move(originatingPrivileges)), needsMerge(needsMerge) { while (authenticatedUsersIter.more()) { @@ -121,7 +99,6 @@ struct ClientCursorParams { const repl::ReadConcernArgs readConcernArgs; int queryOptions = 0; BSONObj originatingCommandObj; - const LockPolicy lockPolicy; PrivilegeVector originatingPrivileges; const bool needsMerge; }; @@ -272,10 +249,6 @@ public: return StringData(_planSummary); } - ClientCursorParams::LockPolicy lockPolicy() const { - return _lockPolicy; - } - /** * Returns a generic cursor containing diagnostics about this cursor. * The caller must either have this cursor pinned or hold a mutex from the cursor manager. @@ -414,8 +387,6 @@ private: // See the QueryOptions enum in dbclientinterface.h. const int _queryOptions = 0; - const ClientCursorParams::LockPolicy _lockPolicy; - // The value of a flag specified on the originating command which indicates whether the result // of this cursor will be consumed by a merging node (mongos or a mongod selected to perform a // merge). Note that this flag is only set for aggregate() commands, and not for find() diff --git a/src/mongo/db/commands/dbcommands.cpp b/src/mongo/db/commands/dbcommands.cpp index 8d9d6b9d9c7..1fc8ae226a7 100644 --- a/src/mongo/db/commands/dbcommands.cpp +++ b/src/mongo/db/commands/dbcommands.cpp @@ -501,7 +501,8 @@ public: result.append("millis", timer.millis()); return 1; } - exec = InternalPlanner::collectionScan(opCtx, ns, collection, PlanExecutor::NO_YIELD); + exec = InternalPlanner::collectionScan( + opCtx, ns, collection, PlanYieldPolicy::YieldPolicy::NO_YIELD); } else if (min.isEmpty() || max.isEmpty()) { errmsg = "only one of min or max specified"; return false; @@ -531,7 +532,7 @@ public: min, max, BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::NO_YIELD); + PlanYieldPolicy::YieldPolicy::NO_YIELD); } long long avgObjSize = collection->dataSize(opCtx) / numRecords; diff --git a/src/mongo/db/commands/dbcommands_d.cpp b/src/mongo/db/commands/dbcommands_d.cpp index 65f6862b318..de2bf3ed410 100644 --- a/src/mongo/db/commands/dbcommands_d.cpp +++ b/src/mongo/db/commands/dbcommands_d.cpp @@ -281,7 +281,7 @@ public: auto exec = uassertStatusOK(getExecutor(opCtx, coll, std::move(cq), - PlanExecutor::YIELD_MANUAL, + PlanYieldPolicy::YieldPolicy::YIELD_MANUAL, QueryPlannerParams::NO_TABLE_SCAN)); // We need to hold a lock to clean up the PlanExecutor, so make sure we have one when we diff --git a/src/mongo/db/commands/dbhash.cpp b/src/mongo/db/commands/dbhash.cpp index f0a4a68c697..918ec3237ea 100644 --- a/src/mongo/db/commands/dbhash.cpp +++ b/src/mongo/db/commands/dbhash.cpp @@ -359,12 +359,12 @@ private: BSONObj(), BSONObj(), BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, InternalPlanner::FORWARD, InternalPlanner::IXSCAN_FETCH); } else if (collection->isCapped()) { exec = InternalPlanner::collectionScan( - opCtx, nss.ns(), collection, PlanExecutor::NO_YIELD); + opCtx, nss.ns(), collection, PlanYieldPolicy::YieldPolicy::NO_YIELD); } else { LOGV2(20455, "Can't find _id index for namespace: {namespace}", diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp index 69f418b5137..6e3462a30bb 100644 --- a/src/mongo/db/commands/find_and_modify.cpp +++ b/src/mongo/db/commands/find_and_modify.cpp @@ -133,8 +133,9 @@ void makeUpdateRequest(OperationContext* opCtx, requestOut->setMulti(false); requestOut->setExplain(explain); - requestOut->setYieldPolicy(opCtx->inMultiDocumentTransaction() ? PlanExecutor::INTERRUPT_ONLY - : PlanExecutor::YIELD_AUTO); + requestOut->setYieldPolicy(opCtx->inMultiDocumentTransaction() + ? PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY + : PlanYieldPolicy::YieldPolicy::YIELD_AUTO); } void makeDeleteRequest(OperationContext* opCtx, @@ -153,8 +154,9 @@ void makeDeleteRequest(OperationContext* opCtx, requestOut->setReturnDeleted(true); // Always return the old value. requestOut->setIsExplain(explain); - requestOut->setYieldPolicy(opCtx->inMultiDocumentTransaction() ? PlanExecutor::INTERRUPT_ONLY - : PlanExecutor::YIELD_AUTO); + requestOut->setYieldPolicy(opCtx->inMultiDocumentTransaction() + ? PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY + : PlanYieldPolicy::YieldPolicy::YIELD_AUTO); } void appendCommandResponse(const PlanExecutor* exec, diff --git a/src/mongo/db/commands/find_cmd.cpp b/src/mongo/db/commands/find_cmd.cpp index b3e114aae79..d9ead1b5891 100644 --- a/src/mongo/db/commands/find_cmd.cpp +++ b/src/mongo/db/commands/find_cmd.cpp @@ -554,8 +554,6 @@ public: // Set up the cursor for getMore. CursorId cursorId = 0; if (shouldSaveCursor(opCtx, collection, state, exec.get())) { - // Create a ClientCursor containing this plan executor and register it with the - // cursor manager. ClientCursorPin pinnedCursor = CursorManager::get(opCtx)->registerCursor( opCtx, {std::move(exec), @@ -564,7 +562,6 @@ public: opCtx->getWriteConcern(), repl::ReadConcernArgs::get(opCtx), _request.body, - ClientCursorParams::LockPolicy::kLockExternally, {Privilege(ResourcePattern::forExactNamespace(nss), ActionType::find)}, expCtx->needsMerge}); cursorId = pinnedCursor.getCursor()->cursorid(); diff --git a/src/mongo/db/commands/getmore_cmd.cpp b/src/mongo/db/commands/getmore_cmd.cpp index 51510ad895b..6ed119df16d 100644 --- a/src/mongo/db/commands/getmore_cmd.cpp +++ b/src/mongo/db/commands/getmore_cmd.cpp @@ -405,7 +405,8 @@ public: PrepareConflictBehavior::kEnforce); } } - if (cursorPin->lockPolicy() == ClientCursorParams::LockPolicy::kLocksInternally) { + if (cursorPin->getExecutor()->lockPolicy() == + PlanExecutor::LockPolicy::kLocksInternally) { if (!_request.nss.isCollectionlessCursorNamespace()) { statsTracker.emplace( opCtx, @@ -415,8 +416,8 @@ public: CollectionCatalog::get(opCtx).getDatabaseProfileLevel(_request.nss.db())); } } else { - invariant(cursorPin->lockPolicy() == - ClientCursorParams::LockPolicy::kLockExternally); + invariant(cursorPin->getExecutor()->lockPolicy() == + PlanExecutor::LockPolicy::kLockExternally); if (MONGO_unlikely(GetMoreHangBeforeReadLock.shouldFail())) { LOGV2(20477, @@ -638,7 +639,8 @@ public: // would be useful to have this info for an aggregation, but the source PlanExecutor // could be destroyed before we know if we need 'execStats' and we do not want to // generate the stats eagerly for all operations due to cost. - if (cursorPin->lockPolicy() != ClientCursorParams::LockPolicy::kLocksInternally && + if (cursorPin->getExecutor()->lockPolicy() != + PlanExecutor::LockPolicy::kLocksInternally && curOp->shouldDBProfile()) { BSONObjBuilder execStatsBob; Explain::getWinningPlanStats(exec, &execStatsBob); diff --git a/src/mongo/db/commands/index_filter_commands_test.cpp b/src/mongo/db/commands/index_filter_commands_test.cpp index d5e50185954..9728caafce2 100644 --- a/src/mongo/db/commands/index_filter_commands_test.cpp +++ b/src/mongo/db/commands/index_filter_commands_test.cpp @@ -105,16 +105,18 @@ vector<BSONObj> getFilters(const QuerySettings& querySettings) { /** * Utility function to create a PlanRankingDecision */ -std::unique_ptr<PlanRankingDecision> createDecision(size_t numPlans) { - unique_ptr<PlanRankingDecision> why(new PlanRankingDecision()); +std::unique_ptr<plan_ranker::PlanRankingDecision> createDecision(size_t numPlans) { + auto why = std::make_unique<plan_ranker::PlanRankingDecision>(); + std::vector<std::unique_ptr<PlanStageStats>> stats; for (size_t i = 0; i < numPlans; ++i) { CommonStats common("COLLSCAN"); - auto stats = std::make_unique<PlanStageStats>(common, STAGE_COLLSCAN); - stats->specific.reset(new CollectionScanStats()); - why->stats.push_back(std::move(stats)); + auto stat = std::make_unique<PlanStageStats>(common, STAGE_COLLSCAN); + stat->specific.reset(new CollectionScanStats()); + stats.push_back(std::move(stat)); why->scores.push_back(0U); why->candidateOrder.push_back(i); } + why->getStats<PlanStageStats>() = std::move(stats); return why; } diff --git a/src/mongo/db/commands/list_collections.cpp b/src/mongo/db/commands/list_collections.cpp index 7c4517c0744..dc590c7f972 100644 --- a/src/mongo/db/commands/list_collections.cpp +++ b/src/mongo/db/commands/list_collections.cpp @@ -368,7 +368,7 @@ public: std::move(ws), std::move(root), nullptr, - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, cursorNss)); for (long long objCount = 0; objCount < batchSize; objCount++) { @@ -405,7 +405,6 @@ public: opCtx->getWriteConcern(), repl::ReadConcernArgs::get(opCtx), jsobj, - ClientCursorParams::LockPolicy::kLocksInternally, uassertStatusOK(AuthorizationSession::get(opCtx->getClient()) ->checkAuthorizedToListCollections(dbname, jsobj)), false // needsMerge always 'false' for listCollections. diff --git a/src/mongo/db/commands/list_indexes.cpp b/src/mongo/db/commands/list_indexes.cpp index b50fbc89ad6..a8df5c10eab 100644 --- a/src/mongo/db/commands/list_indexes.cpp +++ b/src/mongo/db/commands/list_indexes.cpp @@ -173,8 +173,12 @@ public: root->pushBack(id); } - exec = uassertStatusOK(PlanExecutor::make( - expCtx, std::move(ws), std::move(root), nullptr, PlanExecutor::NO_YIELD, nss)); + exec = uassertStatusOK(PlanExecutor::make(expCtx, + std::move(ws), + std::move(root), + nullptr, + PlanYieldPolicy::YieldPolicy::NO_YIELD, + nss)); for (long long objCount = 0; objCount < batchSize; objCount++) { Document nextDoc; @@ -213,7 +217,6 @@ public: opCtx->getWriteConcern(), repl::ReadConcernArgs::get(opCtx), cmdObj, - ClientCursorParams::LockPolicy::kLocksInternally, {Privilege(ResourcePattern::forExactNamespace(nss), ActionType::listIndexes)}, false // needsMerge always 'false' for listIndexes. }); diff --git a/src/mongo/db/commands/run_aggregate.cpp b/src/mongo/db/commands/run_aggregate.cpp index 3816d31ae52..330d74db102 100644 --- a/src/mongo/db/commands/run_aggregate.cpp +++ b/src/mongo/db/commands/run_aggregate.cpp @@ -86,16 +86,6 @@ using std::unique_ptr; namespace { /** - * Returns true if this PlanExecutor is for a Pipeline. - */ -bool isPipelineExecutor(const PlanExecutor* exec) { - invariant(exec); - auto rootStage = exec->getRootStage(); - return rootStage->stageType() == StageType::STAGE_PIPELINE_PROXY || - rootStage->stageType() == StageType::STAGE_CHANGE_STREAM_PROXY; -} - -/** * If a pipeline is empty (assuming that a $cursor stage hasn't been created yet), it could mean * that we were able to absorb all pipeline stages and pull them into a single PlanExecutor. So, * instead of creating a whole pipeline to do nothing more than forward the results of its cursor @@ -121,6 +111,7 @@ bool canOptimizeAwayPipeline(const Pipeline* pipeline, * and thus will be different from that in 'request'. */ bool handleCursorCommand(OperationContext* opCtx, + boost::intrusive_ptr<ExpressionContext> expCtx, const NamespaceString& nsForCursor, std::vector<ClientCursor*> cursors, const AggregationRequest& request, @@ -211,8 +202,7 @@ bool handleCursorCommand(OperationContext* opCtx, // If adding this object will cause us to exceed the message size limit, then we stash it // for later. - BSONObj next = - exec->getExpCtx()->needsMerge ? nextDoc.toBsonWithMetaData() : nextDoc.toBson(); + BSONObj next = expCtx->needsMerge ? nextDoc.toBsonWithMetaData() : nextDoc.toBson(); if (!FindCommon::haveSpaceForNext(next, objCount, responseBuilder.bytesUsed())) { exec->enqueue(nextDoc); stashedResult = true; @@ -475,8 +465,12 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> createOuterPipelineProxyExe // invalidations. The Pipeline may contain PlanExecutors which *are* yielding // PlanExecutors and which *are* registered with their respective collection's // CursorManager - return uassertStatusOK(PlanExecutor::make( - std::move(expCtx), std::move(ws), std::move(proxy), nullptr, PlanExecutor::NO_YIELD, nss)); + return uassertStatusOK(PlanExecutor::make(std::move(expCtx), + std::move(ws), + std::move(proxy), + nullptr, + PlanYieldPolicy::YieldPolicy::NO_YIELD, + nss)); } } // namespace @@ -710,14 +704,6 @@ Status runAggregate(OperationContext* opCtx, } }); for (auto&& exec : execs) { - // PlanExecutors for pipelines always have a 'kLocksInternally' policy. If this executor is - // not for a pipeline, though, that means the pipeline was optimized away and the - // PlanExecutor will answer the query using the query engine only. Without the - // DocumentSourceCursor to do its locking, an executor needs a 'kLockExternally' policy. - auto lockPolicy = isPipelineExecutor(exec.get()) - ? ClientCursorParams::LockPolicy::kLocksInternally - : ClientCursorParams::LockPolicy::kLockExternally; - ClientCursorParams cursorParams( std::move(exec), origNss, @@ -725,7 +711,6 @@ Status runAggregate(OperationContext* opCtx, opCtx->getWriteConcern(), repl::ReadConcernArgs::get(opCtx), cmdObj, - lockPolicy, privileges, expCtx->needsMerge); if (expCtx->tailableMode == TailableModeEnum::kTailable) { @@ -747,13 +732,13 @@ Status runAggregate(OperationContext* opCtx, // If both explain and cursor are specified, explain wins. if (expCtx->explain) { - auto explainExecutor = pins[0].getCursor()->getExecutor(); + auto explainExecutor = pins[0]->getExecutor(); auto bodyBuilder = result->getBodyBuilder(); - if (isPipelineExecutor(explainExecutor)) { + if (explainExecutor->isPipelineExecutor()) { Explain::explainPipelineExecutor(explainExecutor, *(expCtx->explain), &bodyBuilder); } else { - invariant(pins[0].getCursor()->lockPolicy() == - ClientCursorParams::LockPolicy::kLockExternally); + invariant(pins[0]->getExecutor()->lockPolicy() == + PlanExecutor::LockPolicy::kLockExternally); invariant(!explainExecutor->isDetached()); invariant(explainExecutor->getOpCtx() == opCtx); // The explainStages() function for a non-pipeline executor expects to be called with @@ -768,7 +753,7 @@ Status runAggregate(OperationContext* opCtx, } else { // Cursor must be specified, if explain is not. const bool keepCursor = - handleCursorCommand(opCtx, origNss, std::move(cursors), request, result); + handleCursorCommand(opCtx, expCtx, origNss, std::move(cursors), request, result); if (keepCursor) { cursorFreer.dismiss(); } diff --git a/src/mongo/db/commands/test_commands.cpp b/src/mongo/db/commands/test_commands.cpp index d5b87234cba..397bb7188fd 100644 --- a/src/mongo/db/commands/test_commands.cpp +++ b/src/mongo/db/commands/test_commands.cpp @@ -161,8 +161,11 @@ public: // Scan backwards through the collection to find the document to start truncating from. // We will remove 'n' documents, so start truncating from the (n + 1)th document to the // end. - auto exec = InternalPlanner::collectionScan( - opCtx, fullNs.ns(), collection, PlanExecutor::NO_YIELD, InternalPlanner::BACKWARD); + auto exec = InternalPlanner::collectionScan(opCtx, + fullNs.ns(), + collection, + PlanYieldPolicy::YieldPolicy::NO_YIELD, + InternalPlanner::BACKWARD); for (int i = 0; i < n + 1; ++i) { PlanExecutor::ExecState state = exec->getNext(static_cast<BSONObj*>(nullptr), &end); diff --git a/src/mongo/db/commands/write_commands/write_commands.cpp b/src/mongo/db/commands/write_commands/write_commands.cpp index 6d1ea157c21..77b3294cd41 100644 --- a/src/mongo/db/commands/write_commands/write_commands.cpp +++ b/src/mongo/db/commands/write_commands/write_commands.cpp @@ -418,7 +418,7 @@ private: updateRequest.setRuntimeConstants( _batch.getRuntimeConstants().value_or(Variables::generateRuntimeConstants(opCtx))); updateRequest.setLetParameters(_batch.getLet()); - updateRequest.setYieldPolicy(PlanExecutor::YIELD_AUTO); + updateRequest.setYieldPolicy(PlanYieldPolicy::YieldPolicy::YIELD_AUTO); updateRequest.setExplain(verbosity); const ExtensionsCallbackReal extensionsCallback(opCtx, @@ -502,7 +502,7 @@ private: deleteRequest.setQuery(_batch.getDeletes()[0].getQ()); deleteRequest.setCollation(write_ops::collationOf(_batch.getDeletes()[0])); deleteRequest.setMulti(_batch.getDeletes()[0].getMulti()); - deleteRequest.setYieldPolicy(PlanExecutor::YIELD_AUTO); + deleteRequest.setYieldPolicy(PlanYieldPolicy::YieldPolicy::YIELD_AUTO); deleteRequest.setHint(_batch.getDeletes()[0].getHint()); deleteRequest.setIsExplain(true); diff --git a/src/mongo/db/dbhelpers.cpp b/src/mongo/db/dbhelpers.cpp index 9e9c0ada8f9..f6c1c27b4d6 100644 --- a/src/mongo/db/dbhelpers.cpp +++ b/src/mongo/db/dbhelpers.cpp @@ -108,8 +108,8 @@ RecordId Helpers::findOne(OperationContext* opCtx, unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); size_t options = requireIndex ? QueryPlannerParams::NO_TABLE_SCAN : QueryPlannerParams::DEFAULT; - auto exec = uassertStatusOK( - getExecutor(opCtx, collection, std::move(cq), PlanExecutor::NO_YIELD, options)); + auto exec = uassertStatusOK(getExecutor( + opCtx, collection, std::move(cq), PlanYieldPolicy::YieldPolicy::NO_YIELD, options)); PlanExecutor::ExecState state; BSONObj obj; @@ -185,7 +185,8 @@ bool Helpers::getSingleton(OperationContext* opCtx, const char* ns, BSONObj& res boost::optional<AutoGetOplog> autoOplog; auto collection = getCollectionForRead(opCtx, NamespaceString(ns), autoColl, autoOplog); - auto exec = InternalPlanner::collectionScan(opCtx, ns, collection, PlanExecutor::NO_YIELD); + auto exec = InternalPlanner::collectionScan( + opCtx, ns, collection, PlanYieldPolicy::YieldPolicy::NO_YIELD); PlanExecutor::ExecState state = exec->getNext(&result, nullptr); CurOp::get(opCtx)->done(); @@ -207,7 +208,7 @@ bool Helpers::getLast(OperationContext* opCtx, const char* ns, BSONObj& result) auto collection = getCollectionForRead(opCtx, NamespaceString(ns), autoColl, autoOplog); auto exec = InternalPlanner::collectionScan( - opCtx, ns, collection, PlanExecutor::NO_YIELD, InternalPlanner::BACKWARD); + opCtx, ns, collection, PlanYieldPolicy::YieldPolicy::NO_YIELD, InternalPlanner::BACKWARD); PlanExecutor::ExecState state = exec->getNext(&result, nullptr); // Non-yielding collection scans from InternalPlanner will never error. @@ -246,7 +247,7 @@ void Helpers::upsert(OperationContext* opCtx, request.setUpdateModification(updateMod); request.setUpsert(); request.setFromMigration(fromMigrate); - request.setYieldPolicy(PlanExecutor::NO_YIELD); + request.setYieldPolicy(PlanYieldPolicy::YieldPolicy::NO_YIELD); update(opCtx, context.db(), request); } 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; } diff --git a/src/mongo/db/fts/fts_matcher.h b/src/mongo/db/fts/fts_matcher.h index 660194a9585..1f3de2c91f6 100644 --- a/src/mongo/db/fts/fts_matcher.h +++ b/src/mongo/db/fts/fts_matcher.h @@ -75,6 +75,14 @@ public: */ bool negativePhrasesMatch(const BSONObj& obj) const; + const FTSQueryImpl& query() const { + return _query; + } + + const FTSSpec& spec() const { + return _spec; + } + private: /** * For matching, can we skip the positive term check? This is done as optimization when diff --git a/src/mongo/db/index/haystack_access_method.cpp b/src/mongo/db/index/haystack_access_method.cpp index 60accab342a..45a9b2a9443 100644 --- a/src/mongo/db/index/haystack_access_method.cpp +++ b/src/mongo/db/index/haystack_access_method.cpp @@ -134,7 +134,7 @@ void HaystackAccessMethod::searchCommand(OperationContext* opCtx, key, key, BoundInclusion::kIncludeBothStartAndEndKeys, - PlanExecutor::NO_YIELD); + PlanYieldPolicy::YieldPolicy::NO_YIELD); PlanExecutor::ExecState state; BSONObj obj; RecordId loc; diff --git a/src/mongo/db/matcher/expression.h b/src/mongo/db/matcher/expression.h index 052ab19c7ea..e84e3ee0982 100644 --- a/src/mongo/db/matcher/expression.h +++ b/src/mongo/db/matcher/expression.h @@ -36,6 +36,7 @@ #include "mongo/base/status.h" #include "mongo/bson/bsonobj.h" #include "mongo/bson/bsonobjbuilder.h" +#include "mongo/db/matcher/expression_visitor.h" #include "mongo/db/matcher/match_details.h" #include "mongo/db/matcher/matchable.h" #include "mongo/db/pipeline/dependencies.h" @@ -133,6 +134,75 @@ public: }; /** + * An iterator to walk through the children expressions of the given MatchExpressions. Along + * with the defined 'begin()' and 'end()' functions, which take a reference to a + * MatchExpression, this iterator can be used with a range-based loop. For example, + * + * const MatchExpression* expr = makeSomeExpression(); + * for (const auto& child : *expr) { + * ... + * } + * + * When incrementing the iterator, no checks are made to ensure the iterator does not pass + * beyond the boundary. The caller is responsible to compare the iterator against an iterator + * referring to the past-the-end child in the given expression, which can be obtained using + * the 'mongo::end(*expr)' call. + */ + template <bool IsConst> + class MatchExpressionIterator { + public: + MatchExpressionIterator(const MatchExpression* expr, size_t index) + : _expr(expr), _index(index) {} + + template <bool WasConst, typename = std::enable_if_t<IsConst && !WasConst>> + MatchExpressionIterator(const MatchExpressionIterator<WasConst>& other) + : _expr(other._expr), _index(other._index) {} + + template <bool WasConst, typename = std::enable_if_t<IsConst && !WasConst>> + MatchExpressionIterator& operator=(const MatchExpressionIterator<WasConst>& other) { + _expr = other._expr; + _index = other._index; + return *this; + } + + MatchExpressionIterator& operator++() { + ++_index; + return *this; + } + + MatchExpressionIterator operator++(int) { + const auto ret{*this}; + ++(*this); + return ret; + } + + bool operator==(const MatchExpressionIterator& other) const { + return _expr == other._expr && _index == other._index; + } + + bool operator!=(const MatchExpressionIterator& other) const { + return !(*this == other); + } + + template <bool Const = IsConst> + auto operator*() const -> std::enable_if_t<!Const, MatchExpression*> { + return _expr->getChild(_index); + } + + template <bool Const = IsConst> + auto operator*() const -> std::enable_if_t<Const, const MatchExpression*> { + return _expr->getChild(_index); + } + + private: + const MatchExpression* _expr; + size_t _index; + }; + + using Iterator = MatchExpressionIterator<false>; + using ConstIterator = MatchExpressionIterator<true>; + + /** * Make simplifying changes to the structure of a MatchExpression tree without altering its * semantics. This function may return: * - a pointer to the original, unmodified MatchExpression, @@ -338,6 +408,9 @@ public: return false; } + virtual void acceptVisitor(MatchExpressionMutableVisitor* visitor) = 0; + virtual void acceptVisitor(MatchExpressionConstVisitor* visitor) const = 0; + // // Debug information // @@ -397,4 +470,21 @@ private: MatchType _matchType; std::unique_ptr<TagData> _tagData; }; + +inline MatchExpression::Iterator begin(MatchExpression& expr) { + return {&expr, 0}; +} + +inline MatchExpression::ConstIterator begin(const MatchExpression& expr) { + return {&expr, 0}; +} + +inline MatchExpression::Iterator end(MatchExpression& expr) { + return {&expr, expr.numChildren()}; +} + +inline MatchExpression::ConstIterator end(const MatchExpression& expr) { + return {&expr, expr.numChildren()}; +} + } // namespace mongo diff --git a/src/mongo/db/matcher/expression_always_boolean.h b/src/mongo/db/matcher/expression_always_boolean.h index ff12bcd1da0..6b7aff966fe 100644 --- a/src/mongo/db/matcher/expression_always_boolean.h +++ b/src/mongo/db/matcher/expression_always_boolean.h @@ -109,6 +109,14 @@ public: bool isTriviallyFalse() const final { return true; } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class AlwaysTrueMatchExpression final : public AlwaysBooleanMatchExpression { @@ -128,6 +136,14 @@ public: bool isTriviallyTrue() const final { return true; } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/expression_array.h b/src/mongo/db/matcher/expression_array.h index 329d2aea707..67fe81b8e6a 100644 --- a/src/mongo/db/matcher/expression_array.h +++ b/src/mongo/db/matcher/expression_array.h @@ -109,6 +109,14 @@ public: _sub = std::move(newChild); } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; @@ -156,6 +164,14 @@ public: return _subs[i]; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; @@ -201,6 +217,14 @@ public: return _size; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: virtual ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; diff --git a/src/mongo/db/matcher/expression_expr.h b/src/mongo/db/matcher/expression_expr.h index fb9f7d6d628..4dbd889d131 100644 --- a/src/mongo/db/matcher/expression_expr.h +++ b/src/mongo/db/matcher/expression_expr.h @@ -92,6 +92,14 @@ public: return _expression; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; diff --git a/src/mongo/db/matcher/expression_geo.h b/src/mongo/db/matcher/expression_geo.h index 6d38f7fe320..e7b07449744 100644 --- a/src/mongo/db/matcher/expression_geo.h +++ b/src/mongo/db/matcher/expression_geo.h @@ -32,6 +32,7 @@ #include "mongo/db/geo/geometry_container.h" +#include "mongo/db/geo/geoparser.h" #include "mongo/db/matcher/expression.h" #include "mongo/db/matcher/expression_leaf.h" @@ -107,6 +108,14 @@ public: return *_query; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; @@ -191,6 +200,14 @@ public: return *_query; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; @@ -203,4 +220,66 @@ private: std::shared_ptr<const GeoNearExpression> _query; }; +/** + * 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; + } + + std::unique_ptr<MatchExpression> shallowClone() const final { + MONGO_UNREACHABLE; + return nullptr; + } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + +private: + ExpressionOptimizerFunc getOptimizer() const final { + return [](std::unique_ptr<MatchExpression> expression) { return expression; }; + } + + R2Annulus _annulus; +}; } // namespace mongo diff --git a/src/mongo/db/matcher/expression_internal_expr_eq.h b/src/mongo/db/matcher/expression_internal_expr_eq.h index a93c2c52a58..348994f63b5 100644 --- a/src/mongo/db/matcher/expression_internal_expr_eq.h +++ b/src/mongo/db/matcher/expression_internal_expr_eq.h @@ -69,6 +69,14 @@ public: bool matchesSingleElement(const BSONElement&, MatchDetails*) const final; std::unique_ptr<MatchExpression> shallowClone() const final; + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/expression_leaf.cpp b/src/mongo/db/matcher/expression_leaf.cpp index cce1ef9c5fc..18f850ac1e3 100644 --- a/src/mongo/db/matcher/expression_leaf.cpp +++ b/src/mongo/db/matcher/expression_leaf.cpp @@ -199,11 +199,17 @@ constexpr StringData GTEMatchExpression::kName; const std::set<char> RegexMatchExpression::kValidRegexFlags = {'i', 'm', 's', 'x'}; +std::unique_ptr<pcrecpp::RE> RegexMatchExpression::makeRegex(const std::string& regex, + const std::string& flags) { + return std::make_unique<pcrecpp::RE>(regex.c_str(), + regex_util::flagsToPcreOptions(flags, true)); +} + RegexMatchExpression::RegexMatchExpression(StringData path, const BSONElement& e) : LeafMatchExpression(REGEX, path), _regex(e.regex()), _flags(e.regexFlags()), - _re(new pcrecpp::RE(_regex.c_str(), regex_util::flagsToPcreOptions(_flags, true))) { + _re(makeRegex(_regex, _flags)) { uassert(ErrorCodes::BadValue, "regex not a regex", e.type() == RegEx); _init(); } diff --git a/src/mongo/db/matcher/expression_leaf.h b/src/mongo/db/matcher/expression_leaf.h index ebc69b1230a..1cb2b77e545 100644 --- a/src/mongo/db/matcher/expression_leaf.h +++ b/src/mongo/db/matcher/expression_leaf.h @@ -200,6 +200,14 @@ public: e->setCollator(_collator); return std::move(e); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class LTEMatchExpression final : public ComparisonMatchExpression { @@ -222,6 +230,14 @@ public: e->setCollator(_collator); return std::move(e); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class LTMatchExpression final : public ComparisonMatchExpression { @@ -244,6 +260,14 @@ public: e->setCollator(_collator); return std::move(e); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class GTMatchExpression final : public ComparisonMatchExpression { @@ -266,6 +290,14 @@ public: e->setCollator(_collator); return std::move(e); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class GTEMatchExpression final : public ComparisonMatchExpression { @@ -288,12 +320,23 @@ public: e->setCollator(_collator); return std::move(e); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class RegexMatchExpression : public LeafMatchExpression { public: static const std::set<char> kValidRegexFlags; + static std::unique_ptr<pcrecpp::RE> makeRegex(const std::string& regex, + const std::string& flags); + RegexMatchExpression(StringData path, const BSONElement& e); RegexMatchExpression(StringData path, StringData regex, StringData options); @@ -327,6 +370,14 @@ public: return _flags; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; @@ -367,6 +418,14 @@ public: return _remainder; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; @@ -396,6 +455,14 @@ public: virtual bool equivalent(const MatchExpression* other) const; + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; @@ -450,6 +517,14 @@ public: return _hasEmptyArray; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; @@ -571,6 +646,14 @@ public: } return std::move(bitTestMatchExpression); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class BitsAllClearMatchExpression : public BitTestMatchExpression { @@ -592,6 +675,14 @@ public: } return std::move(bitTestMatchExpression); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class BitsAnySetMatchExpression : public BitTestMatchExpression { @@ -613,6 +704,14 @@ public: } return std::move(bitTestMatchExpression); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class BitsAnyClearMatchExpression : public BitTestMatchExpression { @@ -634,6 +733,14 @@ public: } return std::move(bitTestMatchExpression); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/expression_text.h b/src/mongo/db/matcher/expression_text.h index 377ec2d2b21..2b8dba163e7 100644 --- a/src/mongo/db/matcher/expression_text.h +++ b/src/mongo/db/matcher/expression_text.h @@ -64,6 +64,14 @@ public: std::unique_ptr<MatchExpression> shallowClone() const final; + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: fts::FTSQueryImpl _ftsQuery; }; diff --git a/src/mongo/db/matcher/expression_text_noop.h b/src/mongo/db/matcher/expression_text_noop.h index 18a7f2337cf..322b97f7d22 100644 --- a/src/mongo/db/matcher/expression_text_noop.h +++ b/src/mongo/db/matcher/expression_text_noop.h @@ -48,6 +48,14 @@ public: std::unique_ptr<MatchExpression> shallowClone() const final; + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: fts::FTSQueryNoop _ftsQuery; }; diff --git a/src/mongo/db/matcher/expression_tree.h b/src/mongo/db/matcher/expression_tree.h index ddd372db654..d6d1fec2159 100644 --- a/src/mongo/db/matcher/expression_tree.h +++ b/src/mongo/db/matcher/expression_tree.h @@ -131,6 +131,14 @@ public: virtual void serialize(BSONObjBuilder* out, bool includePath) const; bool isTriviallyTrue() const final; + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class OrMatchExpression : public ListOfMatchExpression { @@ -160,6 +168,14 @@ public: virtual void serialize(BSONObjBuilder* out, bool includePath) const; bool isTriviallyFalse() const final; + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class NorMatchExpression : public ListOfMatchExpression { @@ -187,6 +203,14 @@ public: virtual void debugString(StringBuilder& debug, int indentationLevel = 0) const; virtual void serialize(BSONObjBuilder* out, bool includePath) const; + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class NotMatchExpression final : public MatchExpression { @@ -240,6 +264,14 @@ public: return MatchCategory::kLogical; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: static void serializeNotExpressionToNor(MatchExpression* exp, BSONObjBuilder* out, diff --git a/src/mongo/db/matcher/expression_type.h b/src/mongo/db/matcher/expression_type.h index aa2860b8fd2..02369d01b1d 100644 --- a/src/mongo/db/matcher/expression_type.h +++ b/src/mongo/db/matcher/expression_type.h @@ -140,6 +140,14 @@ public: StringData name() const final { return kName; } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; /** @@ -165,6 +173,14 @@ public: MatchCategory getCategory() const final { return MatchCategory::kOther; } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; class InternalSchemaBinDataSubTypeExpression final : public LeafMatchExpression { @@ -227,6 +243,14 @@ public: return _binDataSubType == realOther->_binDataSubType; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; @@ -287,5 +311,13 @@ public: << " of encrypted binary data (0, 1 and 2 are allowed)"); } } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/expression_visitor.h b/src/mongo/db/matcher/expression_visitor.h new file mode 100644 index 00000000000..237263cd4ba --- /dev/null +++ b/src/mongo/db/matcher/expression_visitor.h @@ -0,0 +1,175 @@ +/** + * 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/tree_walker.h" + +namespace mongo { +class AlwaysFalseMatchExpression; +class AlwaysTrueMatchExpression; +class AndMatchExpression; +class BitsAllClearMatchExpression; +class BitsAllSetMatchExpression; +class BitsAnyClearMatchExpression; +class BitsAnySetMatchExpression; +class ElemMatchObjectMatchExpression; +class ElemMatchValueMatchExpression; +class EqualityMatchExpression; +class ExistsMatchExpression; +class ExprMatchExpression; +class GTEMatchExpression; +class GTMatchExpression; +class GeoMatchExpression; +class GeoNearMatchExpression; +class InMatchExpression; +class InternalExprEqMatchExpression; +class InternalSchemaAllElemMatchFromIndexMatchExpression; +class InternalSchemaAllowedPropertiesMatchExpression; +class InternalSchemaBinDataEncryptedTypeExpression; +class InternalSchemaBinDataSubTypeExpression; +class InternalSchemaCondMatchExpression; +class InternalSchemaEqMatchExpression; +class InternalSchemaFmodMatchExpression; +class InternalSchemaMatchArrayIndexMatchExpression; +class InternalSchemaMaxItemsMatchExpression; +class InternalSchemaMaxLengthMatchExpression; +class InternalSchemaMaxPropertiesMatchExpression; +class InternalSchemaMinItemsMatchExpression; +class InternalSchemaMinLengthMatchExpression; +class InternalSchemaMinPropertiesMatchExpression; +class InternalSchemaObjectMatchExpression; +class InternalSchemaRootDocEqMatchExpression; +class InternalSchemaTypeExpression; +class InternalSchemaUniqueItemsMatchExpression; +class InternalSchemaXorMatchExpression; +class LTEMatchExpression; +class LTMatchExpression; +class ModMatchExpression; +class NorMatchExpression; +class NotMatchExpression; +class OrMatchExpression; +class RegexMatchExpression; +class SizeMatchExpression; +class TextMatchExpression; +class TextNoOpMatchExpression; +class TwoDPtInAnnulusExpression; +class TypeMatchExpression; +class WhereMatchExpression; +class WhereNoOpMatchExpression; + +/** + * Visitor pattern for the MatchExpression tree. + * + * This code is not responsible for traversing the tree, only for performing the double-dispatch. + * + * If the visitor doesn't intend to modify the tree, then the template argument 'IsConst' should be + * set to 'true'. In this case all 'visit()' methods will take a const pointer to a visiting node. + */ +template <bool IsConst = false> +class MatchExpressionVisitor { +public: + virtual ~MatchExpressionVisitor() = default; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, AlwaysFalseMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, AlwaysTrueMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, AndMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, BitsAllClearMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, BitsAllSetMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, BitsAnyClearMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, BitsAnySetMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, ElemMatchObjectMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ElemMatchValueMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, EqualityMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ExistsMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ExprMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, GTEMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, GTMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, GeoMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, GeoNearMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, InMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, InternalExprEqMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaAllElemMatchFromIndexMatchExpression> + expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaAllowedPropertiesMatchExpression> + expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaBinDataEncryptedTypeExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaBinDataSubTypeExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaCondMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaEqMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaFmodMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaMatchArrayIndexMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaMaxItemsMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaMaxLengthMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaMaxPropertiesMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaMinItemsMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaMinLengthMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaMinPropertiesMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaObjectMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaRootDocEqMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, InternalSchemaTypeExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaUniqueItemsMatchExpression> expr) = 0; + virtual void visit( + tree_walker::MaybeConstPtr<IsConst, InternalSchemaXorMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, LTEMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, LTMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ModMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, NorMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, NotMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, OrMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, RegexMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, SizeMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, TextMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, TextNoOpMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, TwoDPtInAnnulusExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, TypeMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, WhereMatchExpression> expr) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, WhereNoOpMatchExpression> expr) = 0; +}; + +using MatchExpressionMutableVisitor = MatchExpressionVisitor<false>; +using MatchExpressionConstVisitor = MatchExpressionVisitor<true>; +} // namespace mongo diff --git a/src/mongo/db/matcher/expression_where.h b/src/mongo/db/matcher/expression_where.h index dfe7ee809ea..1e07222dde7 100644 --- a/src/mongo/db/matcher/expression_where.h +++ b/src/mongo/db/matcher/expression_where.h @@ -44,6 +44,14 @@ public: std::unique_ptr<MatchExpression> shallowClone() const final; + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: std::string _dbName; diff --git a/src/mongo/db/matcher/expression_where_noop.h b/src/mongo/db/matcher/expression_where_noop.h index 962f71af20e..38db74cc868 100644 --- a/src/mongo/db/matcher/expression_where_noop.h +++ b/src/mongo/db/matcher/expression_where_noop.h @@ -46,6 +46,14 @@ public: bool matches(const MatchableDocument* doc, MatchDetails* details = nullptr) const final; std::unique_ptr<MatchExpression> shallowClone() const final; + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h b/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h index c7884c63c33..723c1b68084 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h @@ -82,6 +82,14 @@ public: return _expression->getFilter(); } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h b/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h index 739cee4db96..fbe24e8ced3 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h @@ -157,6 +157,14 @@ public: return _patternProperties[i - 1].second->getFilter(); } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_cond.h b/src/mongo/db/matcher/schema/expression_internal_schema_cond.h index 4efd42bd000..d272968ba5f 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_cond.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_cond.h @@ -72,6 +72,14 @@ public: */ bool matches(const MatchableDocument* doc, MatchDetails* details = nullptr) const final; bool matchesSingleElement(const BSONElement& elem, MatchDetails* details = nullptr) const final; + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_eq.h b/src/mongo/db/matcher/schema/expression_internal_schema_eq.h index 69e18a4fc39..ebe8da7e3bd 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_eq.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_eq.h @@ -68,6 +68,14 @@ public: MONGO_UNREACHABLE; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_fmod.h b/src/mongo/db/matcher/schema/expression_internal_schema_fmod.h index 5e99274e069..8f77d0e23aa 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_fmod.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_fmod.h @@ -65,6 +65,14 @@ public: return _remainder; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h b/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h index a5e79e79e7f..a6c063f9244 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_match_array_index.h @@ -87,6 +87,13 @@ public: return _expression->getFilter(); } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_max_items.h b/src/mongo/db/matcher/schema/expression_internal_schema_max_items.h index e133c67e02e..09b8128fc70 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_max_items.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_max_items.h @@ -56,5 +56,13 @@ public: } return std::move(maxItems); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_max_length.h b/src/mongo/db/matcher/schema/expression_internal_schema_max_length.h index 232335afb1b..71ba6513c9e 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_max_length.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_max_length.h @@ -54,6 +54,14 @@ public: } return std::move(maxLen); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_max_properties.h b/src/mongo/db/matcher/schema/expression_internal_schema_max_properties.h index 3e5e05ea849..6394f962028 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_max_properties.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_max_properties.h @@ -66,5 +66,13 @@ public: } return std::move(maxProperties); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_min_items.h b/src/mongo/db/matcher/schema/expression_internal_schema_min_items.h index 2a5978d0326..1c5ad039a9d 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_min_items.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_min_items.h @@ -56,5 +56,13 @@ public: } return std::move(minItems); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_min_length.h b/src/mongo/db/matcher/schema/expression_internal_schema_min_length.h index 06388abd044..b6162443eaf 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_min_length.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_min_length.h @@ -54,6 +54,14 @@ public: } return std::move(minLen); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_min_properties.h b/src/mongo/db/matcher/schema/expression_internal_schema_min_properties.h index 0e9741a281f..4e41b986934 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_min_properties.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_min_properties.h @@ -66,5 +66,13 @@ public: } return std::move(minProperties); } + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h b/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h index 6dda56a1286..e3b40eb9329 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_object_match.h @@ -70,6 +70,14 @@ public: return MatchCategory::kOther; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_root_doc_eq.h b/src/mongo/db/matcher/schema/expression_internal_schema_root_doc_eq.h index a174eabafd2..805529be8f4 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_root_doc_eq.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_root_doc_eq.h @@ -89,6 +89,14 @@ public: return MatchCategory::kOther; } + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + protected: void _doAddDependencies(DepsTracker* deps) const final { deps->needWholeDocument = true; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h b/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h index 328bf24d6f0..f1708b92cef 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_unique_items.h @@ -80,6 +80,14 @@ public: std::unique_ptr<MatchExpression> shallowClone() const final; + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } + private: ExpressionOptimizerFunc getOptimizer() const final { return [](std::unique_ptr<MatchExpression> expression) { return expression; }; diff --git a/src/mongo/db/matcher/schema/expression_internal_schema_xor.h b/src/mongo/db/matcher/schema/expression_internal_schema_xor.h index f64e86c822c..7be046f503e 100644 --- a/src/mongo/db/matcher/schema/expression_internal_schema_xor.h +++ b/src/mongo/db/matcher/schema/expression_internal_schema_xor.h @@ -61,5 +61,13 @@ public: void debugString(StringBuilder& debug, int indentationLevel = 0) const final; void serialize(BSONObjBuilder* out, bool includePath) const final; + + void acceptVisitor(MatchExpressionMutableVisitor* visitor) final { + visitor->visit(this); + } + + void acceptVisitor(MatchExpressionConstVisitor* visitor) const final { + visitor->visit(this); + } }; } // namespace mongo diff --git a/src/mongo/db/ops/delete_request.idl b/src/mongo/db/ops/delete_request.idl index 9198a7ed2f4..9efa4f9b6ed 100644 --- a/src/mongo/db/ops/delete_request.idl +++ b/src/mongo/db/ops/delete_request.idl @@ -42,9 +42,9 @@ types: yield_policy: bson_serialization_type: string description: "The yielding policy of the plan executor" - cpp_type: "PlanExecutor::YieldPolicy" - serializer: "::mongo::PlanExecutor::serializeYieldPolicy" - deserializer: "mongo::PlanExecutor::parseFromBSON" + cpp_type: "PlanYieldPolicy::YieldPolicy" + serializer: "::mongo::PlanYieldPolicy::serializeYieldPolicy" + deserializer: "mongo::PlanYieldPolicy::parseFromBSON" stmt_id: bson_serialization_type: int description: "" @@ -111,4 +111,4 @@ structs: yieldPolicy: description: "The yielding policy of the plan executor." type: yield_policy - default: PlanExecutor::NO_YIELD + default: PlanYieldPolicy::YieldPolicy::NO_YIELD diff --git a/src/mongo/db/ops/parsed_delete.cpp b/src/mongo/db/ops/parsed_delete.cpp index 86b891a0885..ae531650e7f 100644 --- a/src/mongo/db/ops/parsed_delete.cpp +++ b/src/mongo/db/ops/parsed_delete.cpp @@ -132,8 +132,8 @@ const DeleteRequest* ParsedDelete::getRequest() const { return _request; } -PlanExecutor::YieldPolicy ParsedDelete::yieldPolicy() const { - return _request->getGod() ? PlanExecutor::NO_YIELD : _request->getYieldPolicy(); +PlanYieldPolicy::YieldPolicy ParsedDelete::yieldPolicy() const { + return _request->getGod() ? PlanYieldPolicy::YieldPolicy::NO_YIELD : _request->getYieldPolicy(); } bool ParsedDelete::hasParsedQuery() const { diff --git a/src/mongo/db/ops/parsed_delete.h b/src/mongo/db/ops/parsed_delete.h index 09033065604..ccf6b842884 100644 --- a/src/mongo/db/ops/parsed_delete.h +++ b/src/mongo/db/ops/parsed_delete.h @@ -87,7 +87,7 @@ public: /** * Get the YieldPolicy, adjusted for GodMode. */ - PlanExecutor::YieldPolicy yieldPolicy() const; + PlanYieldPolicy::YieldPolicy yieldPolicy() const; /** * As an optimization, we don't create a canonical query for updates with simple _id diff --git a/src/mongo/db/ops/parsed_update.cpp b/src/mongo/db/ops/parsed_update.cpp index 34458e1f8dd..050b4d77c2f 100644 --- a/src/mongo/db/ops/parsed_update.cpp +++ b/src/mongo/db/ops/parsed_update.cpp @@ -230,8 +230,8 @@ ParsedUpdate::parseArrayFilters(const boost::intrusive_ptr<ExpressionContext>& e return std::move(arrayFiltersOut); } -PlanExecutor::YieldPolicy ParsedUpdate::yieldPolicy() const { - return _request->isGod() ? PlanExecutor::NO_YIELD : _request->getYieldPolicy(); +PlanYieldPolicy::YieldPolicy ParsedUpdate::yieldPolicy() const { + return _request->isGod() ? PlanYieldPolicy::YieldPolicy::NO_YIELD : _request->getYieldPolicy(); } bool ParsedUpdate::hasParsedQuery() const { diff --git a/src/mongo/db/ops/parsed_update.h b/src/mongo/db/ops/parsed_update.h index e9be9312389..137ab2084b9 100644 --- a/src/mongo/db/ops/parsed_update.h +++ b/src/mongo/db/ops/parsed_update.h @@ -32,7 +32,7 @@ #include "mongo/base/status.h" #include "mongo/db/matcher/expression_with_placeholder.h" #include "mongo/db/query/collation/collator_interface.h" -#include "mongo/db/query/plan_executor.h" +#include "mongo/db/query/plan_yield_policy.h" #include "mongo/db/update/update_driver.h" namespace mongo { @@ -105,7 +105,7 @@ public: /** * Get the YieldPolicy, adjusted for GodMode. */ - PlanExecutor::YieldPolicy yieldPolicy() const; + PlanYieldPolicy::YieldPolicy yieldPolicy() const; /** * As an optimization, we don't create a canonical query for updates with simple _id diff --git a/src/mongo/db/ops/update_request.h b/src/mongo/db/ops/update_request.h index d5bbbd4bdcc..a18fe293690 100644 --- a/src/mongo/db/ops/update_request.h +++ b/src/mongo/db/ops/update_request.h @@ -228,11 +228,11 @@ public: return shouldReturnOldDocs() || shouldReturnNewDocs(); } - void setYieldPolicy(PlanExecutor::YieldPolicy yieldPolicy) { + void setYieldPolicy(PlanYieldPolicy::YieldPolicy yieldPolicy) { _yieldPolicy = yieldPolicy; } - PlanExecutor::YieldPolicy getYieldPolicy() const { + PlanYieldPolicy::YieldPolicy getYieldPolicy() const { return _yieldPolicy; } @@ -333,7 +333,7 @@ private: ReturnDocOption _returnDocs = ReturnDocOption::RETURN_NONE; // Whether or not the update should yield. Defaults to NO_YIELD. - PlanExecutor::YieldPolicy _yieldPolicy = PlanExecutor::NO_YIELD; + PlanYieldPolicy::YieldPolicy _yieldPolicy = PlanYieldPolicy::YieldPolicy::NO_YIELD; }; } // namespace mongo diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp index f0abbf0d46d..c193a6fd7f8 100644 --- a/src/mongo/db/ops/write_ops_exec.cpp +++ b/src/mongo/db/ops/write_ops_exec.cpp @@ -738,8 +738,9 @@ static SingleWriteResult performSingleUpdateOpWithDupKeyRetry( request.setLetParameters(std::move(letParams)); } request.setStmtId(stmtId); - request.setYieldPolicy(opCtx->inMultiDocumentTransaction() ? PlanExecutor::INTERRUPT_ONLY - : PlanExecutor::YIELD_AUTO); + request.setYieldPolicy(opCtx->inMultiDocumentTransaction() + ? PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY + : PlanYieldPolicy::YieldPolicy::YIELD_AUTO); size_t numAttempts = 0; while (true) { @@ -870,8 +871,9 @@ static SingleWriteResult performSingleDeleteOp(OperationContext* opCtx, request.setQuery(op.getQ()); request.setCollation(write_ops::collationOf(op)); request.setMulti(op.getMulti()); - request.setYieldPolicy(opCtx->inMultiDocumentTransaction() ? PlanExecutor::INTERRUPT_ONLY - : PlanExecutor::YIELD_AUTO); + request.setYieldPolicy(opCtx->inMultiDocumentTransaction() + ? PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY + : PlanYieldPolicy::YieldPolicy::YIELD_AUTO); request.setStmtId(stmtId); request.setHint(op.getHint()); diff --git a/src/mongo/db/pipeline/document_source_cursor.cpp b/src/mongo/db/pipeline/document_source_cursor.cpp index f26604ce3ca..4aa7ce75d11 100644 --- a/src/mongo/db/pipeline/document_source_cursor.cpp +++ b/src/mongo/db/pipeline/document_source_cursor.cpp @@ -135,9 +135,12 @@ void DocumentSourceCursor::loadBatch() { PlanExecutor::ExecState state; Document resultObj; - AutoGetCollectionForRead autoColl(pExpCtx->opCtx, _exec->nss()); - uassertStatusOK(repl::ReplicationCoordinator::get(pExpCtx->opCtx) - ->checkCanServeReadsFor(pExpCtx->opCtx, _exec->nss(), true)); + boost::optional<AutoGetCollectionForRead> autoColl; + if (_exec->lockPolicy() == PlanExecutor::LockPolicy::kLockExternally) { + autoColl.emplace(pExpCtx->opCtx, _exec->nss()); + uassertStatusOK(repl::ReplicationCoordinator::get(pExpCtx->opCtx) + ->checkCanServeReadsFor(pExpCtx->opCtx, _exec->nss(), true)); + } _exec->restoreState(); @@ -160,8 +163,8 @@ void DocumentSourceCursor::loadBatch() { invariant(state == PlanExecutor::IS_EOF); - // Special case for tailable cursor -- EOF doesn't preclude more results, so keep the - // PlanExecutor alive. + // Special case for tailable cursor -- EOF doesn't preclude more results, so keep + // the PlanExecutor alive. if (pExpCtx->isTailableAwaitData()) { _exec->saveState(); return; diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp index eb20fc65a39..42f95c7746f 100644 --- a/src/mongo/db/pipeline/expression.cpp +++ b/src/mongo/db/pipeline/expression.cpp @@ -2363,11 +2363,14 @@ intrusive_ptr<Expression> ExpressionLet::parse(ExpressionContext* const expCtx, auto& inPtr = children.emplace_back(nullptr); std::vector<boost::intrusive_ptr<Expression>>::size_type index = 0; + std::vector<Variables::Id> orderedVariableIds; for (auto&& varElem : varsObj) { const string varName = varElem.fieldName(); Variables::validateNameForUserWrite(varName); Variables::Id id = vpsSub.defineVariable(varName); + orderedVariableIds.push_back(id); + vars.emplace(id, NameAndExpression{varName, children[index]}); // only has outer vars ++index; } @@ -2375,14 +2378,17 @@ intrusive_ptr<Expression> ExpressionLet::parse(ExpressionContext* const expCtx, // parse "in" inPtr = parseOperand(expCtx, inElem, vpsSub); // has our vars - return new ExpressionLet(expCtx, std::move(vars), std::move(children)); + return new ExpressionLet( + expCtx, std::move(vars), std::move(children), std::move(orderedVariableIds)); } ExpressionLet::ExpressionLet(ExpressionContext* const expCtx, VariableMap&& vars, - std::vector<boost::intrusive_ptr<Expression>> children) + std::vector<boost::intrusive_ptr<Expression>> children, + std::vector<Variables::Id> orderedVariableIds) : Expression(expCtx, std::move(children)), _variables(std::move(vars)), + _orderedVariableIds(std::move(orderedVariableIds)), _subExpression(_children.back()) {} intrusive_ptr<Expression> ExpressionLet::optimize() { diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h index 7d1e4d00e3a..4ce4b2c8a01 100644 --- a/src/mongo/db/pipeline/expression.h +++ b/src/mongo/db/pipeline/expression.h @@ -1360,6 +1360,10 @@ public: return _fieldPath; } + Variables::Id getVariableId() const { + return _variable; + } + auto getFieldPathWithoutCurrentPrefix() const { return _fieldPath.tail(); } @@ -1590,6 +1594,10 @@ public: return visitor->visit(this); } + auto& getOrderedVariableIds() const { + return _orderedVariableIds; + } + auto& getVariableMap() const { return _variables; } @@ -1600,9 +1608,15 @@ protected: private: ExpressionLet(ExpressionContext* const expCtx, VariableMap&& vars, - std::vector<boost::intrusive_ptr<Expression>> children); + std::vector<boost::intrusive_ptr<Expression>> children, + std::vector<Variables::Id> orderedVariableIds); VariableMap _variables; + + // These ids are ordered to match their corresponding _children expressions. + std::vector<Variables::Id> _orderedVariableIds; + + // Reference to the last element in the '_children' list. boost::intrusive_ptr<Expression>& _subExpression; }; diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp index 7474d7b95b4..bece5cefcea 100644 --- a/src/mongo/db/pipeline/pipeline_d.cpp +++ b/src/mongo/db/pipeline/pipeline_d.cpp @@ -170,7 +170,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> createRandomCursorEx } return PlanExecutor::make( - expCtx, std::move(ws), std::move(root), coll, PlanExecutor::YIELD_AUTO); + expCtx, std::move(ws), std::move(root), coll, PlanYieldPolicy::YieldPolicy::YIELD_AUTO); } StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> attemptToGetExecutor( diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript index a7e610f3a19..4525a35a990 100644 --- a/src/mongo/db/query/SConscript +++ b/src/mongo/db/query/SConscript @@ -43,6 +43,7 @@ env.Library( "$BUILD_DIR/mongo/db/bson/dotted_path_support", '$BUILD_DIR/mongo/db/commands/server_status_core', "$BUILD_DIR/mongo/db/exec/projection_executor", + "$BUILD_DIR/mongo/db/exec/sbe/query_sbe_plan_stats", "$BUILD_DIR/mongo/db/index/expression_params", "$BUILD_DIR/mongo/db/index/key_generator", "$BUILD_DIR/mongo/db/index_names", @@ -245,6 +246,17 @@ env.Library( ], ) +env.Library( + target="plan_yield_policy", + source=[ + "plan_yield_policy.cpp", + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/util/elapsed_tracker', + ], + ) + env.CppUnitTest( target="db_query_test", source=[ diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp index 38efced1c0c..8afbb78ddb1 100644 --- a/src/mongo/db/query/canonical_query.cpp +++ b/src/mongo/db/query/canonical_query.cpp @@ -33,6 +33,7 @@ #include "mongo/db/query/canonical_query.h" +#include "mongo/db/catalog/collection.h" #include "mongo/db/jsobj.h" #include "mongo/db/matcher/expression_array.h" #include "mongo/db/namespace_string.h" diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/classic_stage_builder.cpp index 87a1987aa2e..a8ef54c2abe 100644 --- a/src/mongo/db/query/stage_builder.cpp +++ b/src/mongo/db/query/classic_stage_builder.cpp @@ -31,7 +31,7 @@ #include "mongo/platform/basic.h" -#include "mongo/db/query/stage_builder.h" +#include "mongo/db/query/classic_stage_builder.h" #include <memory> @@ -64,17 +64,12 @@ #include "mongo/db/storage/oplog_hack.h" #include "mongo/logv2/log.h" -namespace mongo { - +namespace mongo::stage_builder { // Returns a non-null pointer to the root of a plan tree, or a non-OK status if the PlanStage tree // could not be constructed. -std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, - const Collection* collection, - const CanonicalQuery& cq, - const QuerySolution& qsol, - const QuerySolutionNode* root, - WorkingSet* ws) { - auto* const expCtx = cq.getExpCtxRaw(); +std::unique_ptr<PlanStage> ClassicStageBuilder::build(const QuerySolutionNode* root) { + auto* const expCtx = _cq.getExpCtxRaw(); + switch (root->getType()) { case STAGE_COLLSCAN: { const CollectionScanNode* csn = static_cast<const CollectionScanNode*>(root); @@ -90,17 +85,17 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, params.resumeAfterRecordId = csn->resumeAfterRecordId; params.stopApplyingFilterAfterFirstMatch = csn->stopApplyingFilterAfterFirstMatch; return std::make_unique<CollectionScan>( - expCtx, collection, params, ws, csn->filter.get()); + expCtx, _collection, params, _ws, csn->filter.get()); } case STAGE_IXSCAN: { const IndexScanNode* ixn = static_cast<const IndexScanNode*>(root); - invariant(collection); - auto descriptor = collection->getIndexCatalog()->findIndexByName( - opCtx, ixn->index.identifier.catalogName); + invariant(_collection); + auto descriptor = _collection->getIndexCatalog()->findIndexByName( + _opCtx, ixn->index.identifier.catalogName); invariant(descriptor, - str::stream() << "Namespace: " << collection->ns() - << ", CanonicalQuery: " << cq.toStringShort() + str::stream() << "Namespace: " << _collection->ns() + << ", CanonicalQuery: " << _cq.toStringShort() << ", IndexEntry: " << ixn->index.toString()); // We use the node's internal name, keyPattern and multikey details here. For $** @@ -115,21 +110,21 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, params.addKeyMetadata = ixn->addKeyMetadata; params.shouldDedup = ixn->shouldDedup; return std::make_unique<IndexScan>( - expCtx, collection, std::move(params), ws, ixn->filter.get()); + expCtx, _collection, std::move(params), _ws, ixn->filter.get()); } case STAGE_FETCH: { const FetchNode* fn = static_cast<const FetchNode*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, fn->children[0], ws); + auto childStage = build(fn->children[0]); return std::make_unique<FetchStage>( - expCtx, ws, std::move(childStage), fn->filter.get(), collection); + expCtx, _ws, std::move(childStage), fn->filter.get(), _collection); } case STAGE_SORT_DEFAULT: { auto snDefault = static_cast<const SortNodeDefault*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, snDefault->children[0], ws); + auto childStage = build(snDefault->children[0]); return std::make_unique<SortStageDefault>( - cq.getExpCtx(), - ws, - SortPattern{snDefault->pattern, cq.getExpCtx()}, + _cq.getExpCtx(), + _ws, + SortPattern{snDefault->pattern, _cq.getExpCtx()}, snDefault->limit, internalQueryMaxBlockingSortMemoryUsageBytes.load(), snDefault->addSortKeyMetadata, @@ -137,11 +132,11 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, } case STAGE_SORT_SIMPLE: { auto snSimple = static_cast<const SortNodeSimple*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, snSimple->children[0], ws); + auto childStage = build(snSimple->children[0]); return std::make_unique<SortStageSimple>( - cq.getExpCtx(), - ws, - SortPattern{snSimple->pattern, cq.getExpCtx()}, + _cq.getExpCtx(), + _ws, + SortPattern{snSimple->pattern, _cq.getExpCtx()}, snSimple->limit, internalQueryMaxBlockingSortMemoryUsageBytes.load(), snSimple->addSortKeyMetadata, @@ -149,78 +144,77 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, } case STAGE_SORT_KEY_GENERATOR: { const SortKeyGeneratorNode* keyGenNode = static_cast<const SortKeyGeneratorNode*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, keyGenNode->children[0], ws); + auto childStage = build(keyGenNode->children[0]); return std::make_unique<SortKeyGeneratorStage>( - cq.getExpCtx(), std::move(childStage), ws, keyGenNode->sortSpec); + _cq.getExpCtx(), std::move(childStage), _ws, keyGenNode->sortSpec); } case STAGE_RETURN_KEY: { auto returnKeyNode = static_cast<const ReturnKeyNode*>(root); - auto childStage = - buildStages(opCtx, collection, cq, qsol, returnKeyNode->children[0], ws); + auto childStage = build(returnKeyNode->children[0]); return std::make_unique<ReturnKeyStage>( - expCtx, std::move(returnKeyNode->sortKeyMetaFields), ws, std::move(childStage)); + expCtx, std::move(returnKeyNode->sortKeyMetaFields), _ws, std::move(childStage)); } case STAGE_PROJECTION_DEFAULT: { auto pn = static_cast<const ProjectionNodeDefault*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, pn->children[0], ws); - return std::make_unique<ProjectionStageDefault>(cq.getExpCtx(), - cq.getQueryRequest().getProj(), - cq.getProj(), - ws, + auto childStage = build(pn->children[0]); + return std::make_unique<ProjectionStageDefault>(_cq.getExpCtx(), + _cq.getQueryRequest().getProj(), + _cq.getProj(), + _ws, std::move(childStage)); } case STAGE_PROJECTION_COVERED: { auto pn = static_cast<const ProjectionNodeCovered*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, pn->children[0], ws); - return std::make_unique<ProjectionStageCovered>(cq.getExpCtxRaw(), - cq.getQueryRequest().getProj(), - cq.getProj(), - ws, + auto childStage = build(pn->children[0]); + return std::make_unique<ProjectionStageCovered>(_cq.getExpCtxRaw(), + _cq.getQueryRequest().getProj(), + _cq.getProj(), + _ws, std::move(childStage), pn->coveredKeyObj); } case STAGE_PROJECTION_SIMPLE: { auto pn = static_cast<const ProjectionNodeSimple*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, pn->children[0], ws); - return std::make_unique<ProjectionStageSimple>(cq.getExpCtxRaw(), - cq.getQueryRequest().getProj(), - cq.getProj(), - ws, + auto childStage = build(pn->children[0]); + return std::make_unique<ProjectionStageSimple>(_cq.getExpCtxRaw(), + _cq.getQueryRequest().getProj(), + _cq.getProj(), + _ws, std::move(childStage)); } case STAGE_LIMIT: { const LimitNode* ln = static_cast<const LimitNode*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, ln->children[0], ws); - return std::make_unique<LimitStage>(expCtx, ln->limit, ws, std::move(childStage)); + auto childStage = build(ln->children[0]); + return std::make_unique<LimitStage>(expCtx, ln->limit, _ws, std::move(childStage)); } case STAGE_SKIP: { const SkipNode* sn = static_cast<const SkipNode*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, sn->children[0], ws); - return std::make_unique<SkipStage>(expCtx, sn->skip, ws, std::move(childStage)); + auto childStage = build(sn->children[0]); + return std::make_unique<SkipStage>(expCtx, sn->skip, _ws, std::move(childStage)); } case STAGE_AND_HASH: { const AndHashNode* ahn = static_cast<const AndHashNode*>(root); - auto ret = std::make_unique<AndHashStage>(expCtx, ws); + auto ret = std::make_unique<AndHashStage>(expCtx, _ws); for (size_t i = 0; i < ahn->children.size(); ++i) { - auto childStage = buildStages(opCtx, collection, cq, qsol, ahn->children[i], ws); + auto childStage = build(ahn->children[i]); ret->addChild(std::move(childStage)); } return ret; } case STAGE_OR: { const OrNode* orn = static_cast<const OrNode*>(root); - auto ret = std::make_unique<OrStage>(expCtx, ws, orn->dedup, orn->filter.get()); + auto ret = std::make_unique<OrStage>(expCtx, _ws, orn->dedup, orn->filter.get()); for (size_t i = 0; i < orn->children.size(); ++i) { - auto childStage = buildStages(opCtx, collection, cq, qsol, orn->children[i], ws); + auto childStage = build(orn->children[i]); ret->addChild(std::move(childStage)); } return ret; } case STAGE_AND_SORTED: { const AndSortedNode* asn = static_cast<const AndSortedNode*>(root); - auto ret = std::make_unique<AndSortedStage>(expCtx, ws); + auto ret = std::make_unique<AndSortedStage>(expCtx, _ws); for (size_t i = 0; i < asn->children.size(); ++i) { - auto childStage = buildStages(opCtx, collection, cq, qsol, asn->children[i], ws); + auto childStage = build(asn->children[i]); ret->addChild(std::move(childStage)); } return ret; @@ -230,10 +224,10 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, MergeSortStageParams params; params.dedup = msn->dedup; params.pattern = msn->sort; - params.collator = cq.getCollator(); - auto ret = std::make_unique<MergeSortStage>(expCtx, params, ws); + params.collator = _cq.getCollator(); + auto ret = std::make_unique<MergeSortStage>(expCtx, params, _ws); for (size_t i = 0; i < msn->children.size(); ++i) { - auto childStage = buildStages(opCtx, collection, cq, qsol, msn->children[i], ws); + auto childStage = build(msn->children[i]); ret->addChild(std::move(childStage)); } return ret; @@ -248,12 +242,12 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, params.addPointMeta = node->addPointMeta; params.addDistMeta = node->addDistMeta; - invariant(collection); - const IndexDescriptor* twoDIndex = collection->getIndexCatalog()->findIndexByName( - opCtx, node->index.identifier.catalogName); + invariant(_collection); + const IndexDescriptor* twoDIndex = _collection->getIndexCatalog()->findIndexByName( + _opCtx, node->index.identifier.catalogName); invariant(twoDIndex); - return std::make_unique<GeoNear2DStage>(params, expCtx, ws, collection, twoDIndex); + return std::make_unique<GeoNear2DStage>(params, expCtx, _ws, _collection, twoDIndex); } case STAGE_GEO_NEAR_2DSPHERE: { const GeoNear2DSphereNode* node = static_cast<const GeoNear2DSphereNode*>(root); @@ -265,21 +259,22 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, params.addPointMeta = node->addPointMeta; params.addDistMeta = node->addDistMeta; - invariant(collection); - const IndexDescriptor* s2Index = collection->getIndexCatalog()->findIndexByName( - opCtx, node->index.identifier.catalogName); + invariant(_collection); + const IndexDescriptor* s2Index = _collection->getIndexCatalog()->findIndexByName( + _opCtx, node->index.identifier.catalogName); invariant(s2Index); - return std::make_unique<GeoNear2DSphereStage>(params, expCtx, ws, collection, s2Index); + return std::make_unique<GeoNear2DSphereStage>( + params, expCtx, _ws, _collection, s2Index); } case STAGE_TEXT: { const TextNode* node = static_cast<const TextNode*>(root); - invariant(collection); - const IndexDescriptor* desc = collection->getIndexCatalog()->findIndexByName( - opCtx, node->index.identifier.catalogName); + invariant(_collection); + const IndexDescriptor* desc = _collection->getIndexCatalog()->findIndexByName( + _opCtx, node->index.identifier.catalogName); invariant(desc); const FTSAccessMethod* fam = static_cast<const FTSAccessMethod*>( - collection->getIndexCatalog()->getEntry(desc)->accessMethod()); + _collection->getIndexCatalog()->getEntry(desc)->accessMethod()); invariant(fam); TextStageParams params(fam->getSpec()); @@ -289,27 +284,28 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, // practice, this means that it is illegal to use the StageBuilder on a QuerySolution // created by planning a query that contains "no-op" expressions. params.query = static_cast<FTSQueryImpl&>(*node->ftsQuery); - params.wantTextScore = cq.metadataDeps()[DocumentMetadataFields::kTextScore]; - return std::make_unique<TextStage>(expCtx, collection, params, ws, node->filter.get()); + params.wantTextScore = _cq.metadataDeps()[DocumentMetadataFields::kTextScore]; + return std::make_unique<TextStage>( + expCtx, _collection, params, _ws, node->filter.get()); } case STAGE_SHARDING_FILTER: { const ShardingFilterNode* fn = static_cast<const ShardingFilterNode*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, fn->children[0], ws); + auto childStage = build(fn->children[0]); - auto css = CollectionShardingState::get(opCtx, collection->ns()); + auto css = CollectionShardingState::get(_opCtx, _collection->ns()); return std::make_unique<ShardFilterStage>( expCtx, css->getOwnershipFilter( - opCtx, CollectionShardingState::OrphanCleanupPolicy::kDisallowOrphanCleanup), - ws, + _opCtx, CollectionShardingState::OrphanCleanupPolicy::kDisallowOrphanCleanup), + _ws, std::move(childStage)); } case STAGE_DISTINCT_SCAN: { const DistinctNode* dn = static_cast<const DistinctNode*>(root); - invariant(collection); - auto descriptor = collection->getIndexCatalog()->findIndexByName( - opCtx, dn->index.identifier.catalogName); + invariant(_collection); + auto descriptor = _collection->getIndexCatalog()->findIndexByName( + _opCtx, dn->index.identifier.catalogName); invariant(descriptor); // We use the node's internal name, keyPattern and multikey details here. For $** @@ -323,14 +319,14 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, params.scanDirection = dn->direction; params.bounds = dn->bounds; params.fieldNo = dn->fieldNo; - return std::make_unique<DistinctScan>(expCtx, collection, std::move(params), ws); + return std::make_unique<DistinctScan>(expCtx, _collection, std::move(params), _ws); } case STAGE_COUNT_SCAN: { const CountScanNode* csn = static_cast<const CountScanNode*>(root); - invariant(collection); - auto descriptor = collection->getIndexCatalog()->findIndexByName( - opCtx, csn->index.identifier.catalogName); + invariant(_collection); + auto descriptor = _collection->getIndexCatalog()->findIndexByName( + _opCtx, csn->index.identifier.catalogName); invariant(descriptor); // We use the node's internal name, keyPattern and multikey details here. For $** @@ -345,13 +341,13 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, params.startKeyInclusive = csn->startKeyInclusive; params.endKey = csn->endKey; params.endKeyInclusive = csn->endKeyInclusive; - return std::make_unique<CountScan>(expCtx, collection, std::move(params), ws); + return std::make_unique<CountScan>(expCtx, _collection, std::move(params), _ws); } case STAGE_ENSURE_SORTED: { const EnsureSortedNode* esn = static_cast<const EnsureSortedNode*>(root); - auto childStage = buildStages(opCtx, collection, cq, qsol, esn->children[0], ws); + auto childStage = build(esn->children[0]); return std::make_unique<EnsureSortedStage>( - expCtx, esn->pattern, ws, std::move(childStage)); + expCtx, esn->pattern, _ws, std::move(childStage)); } case STAGE_CACHED_PLAN: case STAGE_CHANGE_STREAM_PROXY: @@ -380,23 +376,4 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, MONGO_UNREACHABLE; } - -std::unique_ptr<PlanStage> StageBuilder::build(OperationContext* opCtx, - const Collection* collection, - const CanonicalQuery& cq, - const QuerySolution& solution, - WorkingSet* wsIn) { - // Only QuerySolutions derived from queries parsed with context, or QuerySolutions derived from - // queries that disallow extensions, can be properly executed. If the query does not have - // $text/$where context (and $text/$where are allowed), then no attempt should be made to - // execute the query. - invariant(!cq.canHaveNoopMatchNodes()); - - invariant(wsIn); - invariant(solution.root); - - QuerySolutionNode* solutionNode = solution.root.get(); - return buildStages(opCtx, collection, cq, solution, solutionNode, wsIn); -} - -} // namespace mongo +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/classic_stage_builder.h b/src/mongo/db/query/classic_stage_builder.h new file mode 100644 index 00000000000..bdf68d4677b --- /dev/null +++ b/src/mongo/db/query/classic_stage_builder.h @@ -0,0 +1,53 @@ +/** + * 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/plan_stage.h" +#include "mongo/db/query/stage_builder.h" + +namespace mongo::stage_builder { +/** + * A stage builder which builds an executable tree using classic PlanStages. + */ +class ClassicStageBuilder : public StageBuilder<PlanStage> { +public: + ClassicStageBuilder(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + const QuerySolution& solution, + WorkingSet* ws) + : StageBuilder<PlanStage>{opCtx, collection, cq, solution}, _ws{ws} {} + + std::unique_ptr<PlanStage> build(const QuerySolutionNode* root) final; + +private: + WorkingSet* _ws; +}; +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/explain.cpp b/src/mongo/db/query/explain.cpp index 96fae9342db..4179c65db85 100644 --- a/src/mongo/db/query/explain.cpp +++ b/src/mongo/db/query/explain.cpp @@ -50,6 +50,7 @@ #include "mongo/db/query/explain_common.h" #include "mongo/db/query/get_executor.h" #include "mongo/db/query/plan_executor.h" +#include "mongo/db/query/plan_executor_sbe.h" #include "mongo/db/query/plan_summary_stats.h" #include "mongo/db/query/query_planner.h" #include "mongo/db/query/query_settings.h" @@ -626,6 +627,12 @@ BSONObj Explain::statsToBSON(const PlanStageStats& stats, ExplainOptions::Verbos return bob.obj(); } +BSONObj Explain::statsToBSON(const sbe::PlanStageStats& stats, + ExplainOptions::Verbosity verbosity) { + BSONObjBuilder bob; + return bob.obj(); +} + // static void Explain::statsToBSON(const PlanStageStats& stats, BSONObjBuilder* bob, @@ -636,7 +643,9 @@ void Explain::statsToBSON(const PlanStageStats& stats, // static BSONObj Explain::getWinningPlanStats(const PlanExecutor* exec) { BSONObjBuilder bob; - getWinningPlanStats(exec, &bob); + if (!dynamic_cast<const PlanExecutorSBE*>(exec)) { + getWinningPlanStats(exec, &bob); + } return bob.obj(); } @@ -882,6 +891,10 @@ void Explain::explainStages(PlanExecutor* exec, ExplainOptions::Verbosity verbosity, BSONObj extraInfo, BSONObjBuilder* out) { + uassert(4822877, + "Explain facility is not supported for SBE plans", + !dynamic_cast<PlanExecutorSBE*>(exec)); + auto winningPlanTrialStats = Explain::getWinningPlanTrialStats(exec); Status executePlanStatus = Status::OK(); @@ -915,6 +928,11 @@ void Explain::explainStages(PlanExecutor* exec, // static std::string Explain::getPlanSummary(const PlanExecutor* exec) { + // TODO: Handle planSummary when SBE is enabled. + if (dynamic_cast<const PlanExecutorSBE*>(exec)) { + return "unsupported"; + } + return getPlanSummary(exec->getRootStage()); } @@ -950,11 +968,21 @@ std::string Explain::getPlanSummary(const PlanStage* root) { } // static +std::string Explain::getPlanSummary(const sbe::PlanStage* root) { + // TODO: Handle 'planSummary' when SBE is enabled. + return "unsupported"; +} + +// static void Explain::getSummaryStats(const PlanExecutor& exec, PlanSummaryStats* statsOut) { invariant(nullptr != statsOut); - PlanStage* root = exec.getRootStage(); + // TODO: Handle 'getSummaryStats' when SBE is enabled. + if (dynamic_cast<const PlanExecutorSBE*>(&exec)) { + return; + } + PlanStage* root = exec.getRootStage(); if (root->stageType() == STAGE_PIPELINE_PROXY || root->stageType() == STAGE_CHANGE_STREAM_PROXY) { auto pipelineProxy = static_cast<PipelineProxyStage*>(root); @@ -1055,14 +1083,15 @@ void Explain::planCacheEntryToBSON(const PlanCacheEntry& entry, BSONObjBuilder* out->append("works", static_cast<long long>(entry.works)); BSONObjBuilder cachedPlanBob(out->subobjStart("cachedPlan")); - Explain::statsToBSON( - *entry.decision->stats[0], &cachedPlanBob, ExplainOptions::Verbosity::kQueryPlanner); + Explain::statsToBSON(*(entry.decision->getStats<PlanStageStats>()[0]), + &cachedPlanBob, + ExplainOptions::Verbosity::kQueryPlanner); cachedPlanBob.doneFast(); out->append("timeOfCreation", entry.timeOfCreation); BSONArrayBuilder creationBuilder(out->subarrayStart("creationExecStats")); - for (auto&& stat : entry.decision->stats) { + for (auto&& stat : entry.decision->getStats<PlanStageStats>()) { BSONObjBuilder planBob(creationBuilder.subobjStart()); Explain::generateSinglePlanExecutionInfo( stat.get(), ExplainOptions::Verbosity::kExecAllPlans, boost::none, &planBob); diff --git a/src/mongo/db/query/explain.h b/src/mongo/db/query/explain.h index 2fff62ac1d2..2cc16c85b0a 100644 --- a/src/mongo/db/query/explain.h +++ b/src/mongo/db/query/explain.h @@ -128,6 +128,9 @@ public: static BSONObj statsToBSON( const PlanStageStats& stats, ExplainOptions::Verbosity verbosity = ExplainOptions::Verbosity::kExecStats); + static BSONObj statsToBSON( + const sbe::PlanStageStats& stats, + ExplainOptions::Verbosity verbosity = ExplainOptions::Verbosity::kExecStats); /** * This version of stats tree to BSON conversion returns the result through the @@ -146,6 +149,7 @@ public: */ static std::string getPlanSummary(const PlanExecutor* exec); static std::string getPlanSummary(const PlanStage* root); + static std::string getPlanSummary(const sbe::PlanStage* root); /** * Fills out 'statsOut' with summary stats using the execution tree contained diff --git a/src/mongo/db/query/find.cpp b/src/mongo/db/query/find.cpp index 602e525cb9c..49169ef4ddc 100644 --- a/src/mongo/db/query/find.cpp +++ b/src/mongo/db/query/find.cpp @@ -270,7 +270,7 @@ Message getMore(OperationContext* opCtx, opCtx->setExhaust(cursorPin->queryOptions() & QueryOption_Exhaust); - if (cursorPin->lockPolicy() == ClientCursorParams::LockPolicy::kLocksInternally) { + if (cursorPin->getExecutor()->lockPolicy() == PlanExecutor::LockPolicy::kLocksInternally) { if (!nss.isCollectionlessCursorNamespace()) { AutoGetDb autoDb(opCtx, nss.db(), MODE_IS); statsTracker.emplace(opCtx, @@ -497,7 +497,7 @@ Message getMore(OperationContext* opCtx, // info for an aggregation, but the source PlanExecutor could be destroyed before we know if we // need 'execStats' and we do not want to generate the stats eagerly for all operations due to // cost. - if (cursorPin->lockPolicy() != ClientCursorParams::LockPolicy::kLocksInternally && + if (cursorPin->getExecutor()->lockPolicy() != PlanExecutor::LockPolicy::kLocksInternally && curOp.shouldDBProfile()) { BSONObjBuilder execStatsBob; Explain::getWinningPlanStats(exec, &execStatsBob); @@ -756,7 +756,6 @@ bool runQuery(OperationContext* opCtx, opCtx->getWriteConcern(), readConcernArgs, upconvertedQuery, - ClientCursorParams::LockPolicy::kLockExternally, {Privilege(ResourcePattern::forExactNamespace(nss), ActionType::find)}, false // needsMerge always 'false' for find(). }); diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index 42314ec67aa..50d1ef4d823 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -51,6 +51,8 @@ #include "mongo/db/exec/projection_executor_utils.h" #include "mongo/db/exec/record_store_fast_count.h" #include "mongo/db/exec/return_key.h" +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" #include "mongo/db/exec/shard_filter.h" #include "mongo/db/exec/sort_key_generator.h" #include "mongo/db/exec/subplan.h" @@ -79,7 +81,11 @@ #include "mongo/db/query/query_planner.h" #include "mongo/db/query/query_planner_common.h" #include "mongo/db/query/query_settings.h" -#include "mongo/db/query/stage_builder.h" +#include "mongo/db/query/sbe_cached_solution_planner.h" +#include "mongo/db/query/sbe_multi_planner.h" +#include "mongo/db/query/sbe_sub_planner.h" +#include "mongo/db/query/stage_builder_util.h" +#include "mongo/db/query/util/make_data_structure.h" #include "mongo/db/repl/optime.h" #include "mongo/db/repl/replication_coordinator.h" #include "mongo/db/s/collection_sharding_state.h" @@ -93,10 +99,6 @@ namespace mongo { -using std::string; -using std::unique_ptr; -using std::vector; - boost::intrusive_ptr<ExpressionContext> makeExpressionContextForGetExecutor( OperationContext* opCtx, const BSONObj& requestCollation, const NamespaceString& nss) { invariant(opCtx); @@ -137,6 +139,18 @@ namespace { namespace wcp = ::mongo::wildcard_planning; // The body is below in the "count hack" section but getExecutor calls it. bool turnIxscanIntoCount(QuerySolution* soln); + +/** + * Returns 'true' if 'query' on the given 'collection' can be answered using a special IDHACK plan. + */ +bool isIdHackEligibleQuery(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()); +} } // namespace bool isAnyComponentOfPathMultikey(const BSONObj& indexKeyPattern, @@ -236,7 +250,7 @@ IndexEntry indexEntryFromIndexCatalogEntry(OperationContext* opCtx, void applyIndexFilters(Collection* collection, const CanonicalQuery& canonicalQuery, QueryPlannerParams* plannerParams) { - if (!IDHackStage::supportsQuery(collection, canonicalQuery)) { + if (!isIdHackEligibleQuery(collection, canonicalQuery)) { QuerySettings* querySettings = CollectionQueryInfo::get(collection).getQuerySettings(); const auto key = canonicalQuery.encodeKey(); @@ -343,301 +357,739 @@ bool shouldWaitForOplogVisibility(OperationContext* opCtx, } namespace { +/** + * A base class to hold the result returned by PrepareExecutionHelper::prepare call. + */ +template <typename PlanStageType> +class BasePrepareExecutionResult { +public: + BasePrepareExecutionResult() = default; + virtual ~BasePrepareExecutionResult() = default; + + /** + * Saves the provided PlanStage 'root' and the 'solution' object within this instance. + * + * The exact semantics of this method is defined by a specific subclass. For example, this + * result object can store a single execution tree for the canonical query, even though it may + * have multiple solutions. In this the case the execution tree itself would encapsulate the + * multi planner logic to choose the best plan in runtime. + * + * Or, alternatively, this object can store multiple execution trees (and solutions). In this + * case each 'root' and 'solution' objects passed to this method would have to be stored in an + * internal vector to be later returned to the user of this class, in case runtime planning is + * implemented outside of the execution tree. + */ + virtual void emplace(PlanStageType root, std::unique_ptr<QuerySolution> solution) = 0; + + /** + * Returns a short plan summary describing the shape of the query plan. + * + * Should only be called when this result contains a single execution tree and a query solution. + */ + virtual std::string getPlanSummary() const = 0; +}; + +/** + * A class to hold the result of preparation of the query to be executed using classic engine. This + * result stores and provides the following information: + * - A QuerySolutions for the query. May be null in certain circumstances, where the constructed + * execution tree does not have an associated query solution. + * - A root PlanStage of the constructed execution tree. + */ +class ClassicPrepareExecutionResult final + : public BasePrepareExecutionResult<std::unique_ptr<PlanStage>> { +public: + using BasePrepareExecutionResult::BasePrepareExecutionResult; + + void emplace(std::unique_ptr<PlanStage> root, std::unique_ptr<QuerySolution> solution) final { + invariant(!_root); + invariant(!_solution); + _root = std::move(root); + _solution = std::move(solution); + } + + std::string getPlanSummary() const final { + invariant(_root); + return Explain::getPlanSummary(_root.get()); + } -struct PrepareExecutionResult { - PrepareExecutionResult(unique_ptr<CanonicalQuery> canonicalQuery, - unique_ptr<QuerySolution> querySolution, - unique_ptr<PlanStage> root) - : canonicalQuery(std::move(canonicalQuery)), - querySolution(std::move(querySolution)), - root(std::move(root)) {} - - unique_ptr<CanonicalQuery> canonicalQuery; - unique_ptr<QuerySolution> querySolution; - unique_ptr<PlanStage> root; + std::unique_ptr<PlanStage> root() { + return std::move(_root); + } + + std::unique_ptr<QuerySolution> solution() { + return std::move(_solution); + } + +private: + std::unique_ptr<PlanStage> _root; + std::unique_ptr<QuerySolution> _solution; }; /** - * Build an execution tree for the query described in 'canonicalQuery'. - * - * If an execution tree could be created, then returns a PrepareExecutionResult that wraps: - * - The CanonicalQuery describing the query operation. This may be equal to the original canonical - * query, or may be modified. This will never be null. - * - A QuerySolution, representing the associated query solution. This may be null, in certain - * circumstances where the constructed execution tree does not have an associated query solution. - * - A PlanStage, representing the root of the constructed execution tree. This will never be null. - * - * If an execution tree could not be created, returns an error Status. + * A class to hold the result of preparation of the query to be executed using SBE engine. This + * result stores and provides the following information: + * - A vector of QuerySolutions. Elements of the vector may be null, in certain circumstances + * where the constructed execution tree does not have an associated query solution. + * - A vector of PlanStages, representing the roots of the constructed execution trees (in the + * case when the query has multiple solutions, we may construct an execution tree for each + * solution and pick the best plan after multi-planning). Elements of this vector can never be + * null. The size of this vector must always match the size of 'querySolutions' vector. + * - An optional decisionWorks value, which is populated when a solution was reconstructed from + * the PlanCache, and will hold the number of work cycles taken to decide on a winning plan + * when the plan was first cached. It used to decided whether cached solution runtime planning + * needs to be done or not. + * - A 'needSubplanning' flag indicating that the query contains rooted $or predicate and is + * eligible for runtime sub-planning. */ -StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx, - Collection* collection, - WorkingSet* ws, - unique_ptr<CanonicalQuery> canonicalQuery, - size_t plannerOptions) { - invariant(canonicalQuery); - unique_ptr<PlanStage> root; +class SlotBasedPrepareExecutionResult final + : public BasePrepareExecutionResult< + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData>> { +public: + using QuerySolutionVector = std::vector<std::unique_ptr<QuerySolution>>; + using PlanStageVector = + std::vector<std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData>>; - // This can happen as we're called by internal clients as well. - if (nullptr == collection) { - const string& ns = canonicalQuery->ns(); - LOGV2_DEBUG(20921, - 2, - "Collection {namespace} does not exist. Using EOF plan: {query}", - "Collection does not exist. Using EOF plan", - "namespace"_attr = ns, - "query"_attr = redact(canonicalQuery->toStringShort())); - root = std::make_unique<EOFStage>(canonicalQuery->getExpCtxRaw()); - return PrepareExecutionResult(std::move(canonicalQuery), nullptr, std::move(root)); + using BasePrepareExecutionResult::BasePrepareExecutionResult; + + void emplace(std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root, + std::unique_ptr<QuerySolution> solution) final { + _roots.push_back(std::move(root)); + _solutions.push_back(std::move(solution)); } - // Fill out the planning params. We use these for both cached solutions and non-cached. - QueryPlannerParams plannerParams; - plannerParams.options = plannerOptions; - fillOutPlannerParams(opCtx, collection, canonicalQuery.get(), &plannerParams); + std::string getPlanSummary() const final { + // We can report plan summary only if this result contains a single solution. + invariant(_roots.size() == 1); + invariant(_roots[0].first); + return Explain::getPlanSummary(_roots[0].first.get()); + } - // If the canonical query does not have a user-specified collation and no one has given the - // CanonicalQuery a collation already, set it from the collection default. - if (canonicalQuery->getQueryRequest().getCollation().isEmpty() && - canonicalQuery->getCollator() == nullptr && collection->getDefaultCollator()) { - canonicalQuery->setCollator(collection->getDefaultCollator()->clone()); + PlanStageVector roots() { + return std::move(_roots); } - const IndexDescriptor* descriptor = collection->getIndexCatalog()->findIdIndex(opCtx); + QuerySolutionVector solutions() { + return std::move(_solutions); + } - // If we have an _id index we can use an idhack plan. - if (descriptor && IDHackStage::supportsQuery(collection, *canonicalQuery)) { - LOGV2_DEBUG(20922, - 2, - "Using idhack: {canonicalQuery_Short}", - "canonicalQuery_Short"_attr = redact(canonicalQuery->toStringShort())); + boost::optional<size_t> decisionWorks() const { + return _decisionWorks; + } - root = std::make_unique<IDHackStage>( - canonicalQuery->getExpCtxRaw(), canonicalQuery.get(), ws, collection, descriptor); + bool needsSubplanning() const { + return _needSubplanning; + } + + void setNeedsSubplanning(bool needsSubplanning) { + _needSubplanning = needsSubplanning; + } + + void setDecisionWorks(size_t decisionWorks) { + _decisionWorks = decisionWorks; + } + +private: + QuerySolutionVector _solutions; + PlanStageVector _roots; + boost::optional<size_t> _decisionWorks; + bool _needSubplanning{false}; +}; + +/** + * A helper class to build and prepare a PlanStage tree for execution. This class contains common + * logic to build and prepare an execution tree for the provided canonical query, and also provides + * methods to build various specialized PlanStage trees when we either: + * * Do not build a QuerySolutionNode tree for the input query, and as such do not undergo the + * normal stage builder process. + * * We have a QuerySolutionNode tree (or multiple query solution trees), but must execute some + * custom logic in order to build the final execution tree. + */ +template <typename PlanStageType, typename ResultType> +class PrepareExecutionHelper { +public: + PrepareExecutionHelper(OperationContext* opCtx, + Collection* collection, + CanonicalQuery* cq, + PlanYieldPolicy* yieldPolicy, + size_t plannerOptions) + : _opCtx{opCtx}, + _collection{collection}, + _cq{cq}, + _yieldPolicy{yieldPolicy}, + _plannerOptions{plannerOptions} { + invariant(_cq); + } + + StatusWith<std::unique_ptr<ResultType>> prepare() { + if (nullptr == _collection) { + LOGV2_DEBUG(20921, + 2, + "Collection {namespace} does not exist. Using EOF plan: {query}", + "Collection does not exist. Using EOF plan", + "namespace"_attr = _cq->ns(), + "canonicalQuery"_attr = redact(_cq->toStringShort())); + return buildEofPlan(); + } + + // Fill out the planning params. We use these for both cached solutions and non-cached. + QueryPlannerParams plannerParams; + plannerParams.options = _plannerOptions; + fillOutPlannerParams(_opCtx, _collection, _cq, &plannerParams); + + // If the canonical query does not have a user-specified collation and no one has given the + // CanonicalQuery a collation already, set it from the collection default. + if (_cq->getQueryRequest().getCollation().isEmpty() && _cq->getCollator() == nullptr && + _collection->getDefaultCollator()) { + _cq->setCollator(_collection->getDefaultCollator()->clone()); + } + + const IndexDescriptor* idIndexDesc = _collection->getIndexCatalog()->findIdIndex(_opCtx); + + // If we have an _id index we can use an idhack plan. + if (idIndexDesc && isIdHackEligibleQuery(_collection, *_cq)) { + LOGV2_DEBUG( + 20922, 2, "Using idhack", "canonicalQuery"_attr = redact(_cq->toStringShort())); + // If an IDHACK plan is not supported, we will use the normal plan generation process + // to force an _id index scan plan. If an IDHACK plan was generated we return + // immediately, otherwise we fall through and continue. + if (auto result = buildIdHackPlan(idIndexDesc, &plannerParams)) { + return std::move(result); + } + } + + // Tailable: If the query requests tailable the collection must be capped. + if (_cq->getQueryRequest().isTailable() && !_collection->isCapped()) { + return Status(ErrorCodes::BadValue, + str::stream() << "error processing query: " << _cq->toString() + << " tailable cursor requested on non capped collection"); + } + + // Check that the query should be cached. + if (CollectionQueryInfo::get(_collection).getPlanCache()->shouldCacheQuery(*_cq)) { + // Fill in opDebug information. + const auto planCacheKey = + CollectionQueryInfo::get(_collection).getPlanCache()->computeKey(*_cq); + CurOp::get(_opCtx)->debug().queryHash = + canonical_query_encoder::computeHash(planCacheKey.getStableKeyStringData()); + CurOp::get(_opCtx)->debug().planCacheKey = + canonical_query_encoder::computeHash(planCacheKey.toString()); + + // Try to look up a cached solution for the query. + if (auto cs = CollectionQueryInfo::get(_collection) + .getPlanCache() + ->getCacheEntryIfActive(planCacheKey)) { + // We have a CachedSolution. Have the planner turn it into a QuerySolution. + auto statusWithQs = QueryPlanner::planFromCache(*_cq, plannerParams, *cs); + + if (statusWithQs.isOK()) { + auto querySolution = std::move(statusWithQs.getValue()); + if ((plannerParams.options & QueryPlannerParams::IS_COUNT) && + turnIxscanIntoCount(querySolution.get())) { + LOGV2_DEBUG(20923, + 2, + "Using fast count: {query}", + "Using fast count", + "query"_attr = redact(_cq->toStringShort())); + } + + return buildCachedPlan( + std::move(querySolution), plannerParams, cs->decisionWorks); + } + } + } + + + if (internalQueryPlanOrChildrenIndependently.load() && + SubplanStage::canUseSubplanning(*_cq)) { + LOGV2_DEBUG(20924, + 2, + "Running query as sub-queries: {query}", + "Running query as sub-queries", + "query"_attr = redact(_cq->toStringShort())); + return buildSubPlan(plannerParams); + } + + auto statusWithSolutions = QueryPlanner::plan(*_cq, plannerParams); + if (!statusWithSolutions.isOK()) { + return statusWithSolutions.getStatus().withContext( + str::stream() << "error processing query: " << _cq->toString() + << " planner returned error"); + } + auto solutions = std::move(statusWithSolutions.getValue()); + // The planner should have returned an error status if there are no solutions. + invariant(solutions.size() > 0); + + // See if one of our solutions is a fast count hack in disguise. + if (plannerParams.options & QueryPlannerParams::IS_COUNT) { + for (size_t i = 0; i < solutions.size(); ++i) { + if (turnIxscanIntoCount(solutions[i].get())) { + auto result = makeResult(); + auto root = buildExecutableTree(*solutions[i]); + result->emplace(std::move(root), std::move(solutions[i])); + + LOGV2_DEBUG(20925, + 2, + "Using fast count: {query}, planSummary: {planSummary}", + "Using fast count", + "query"_attr = redact(_cq->toStringShort()), + "planSummary"_attr = result->getPlanSummary()); + return std::move(result); + } + } + } + + if (1 == solutions.size()) { + auto result = makeResult(); + // Only one possible plan. Run it. Build the stages from the solution. + auto root = buildExecutableTree(*solutions[0]); + result->emplace(std::move(root), std::move(solutions[0])); + + LOGV2_DEBUG( + 20926, + 2, + "Only one plan is available; it will be run but will not be cached. {query}, " + "planSummary: {planSummary}", + "Only one plan is available; it will be run but will not be cached", + "query"_attr = redact(_cq->toStringShort()), + "planSummary"_attr = result->getPlanSummary()); + + return std::move(result); + } + + return buildMultiPlan(std::move(solutions), plannerParams); + } + +protected: + /** + * Creates a result instance to be returned to the caller holding the result of the + * prepare() call. + */ + auto makeResult() const { + return std::make_unique<ResultType>(); + } + + /** + * Constructs a PlanStage tree from the given query 'solution'. + */ + virtual PlanStageType buildExecutableTree(const QuerySolution& solution) const = 0; + + /** + * Constructs a special PlanStage tree to return EOF immediately on the first call to getNext. + */ + virtual std::unique_ptr<ResultType> buildEofPlan() = 0; + + /** + * If supported, constructs a special PlanStage tree for fast-path document retrievals via the + * _id index. Otherwise, nullptr should be returned and this helper will fall back to the + * normal plan generation. + */ + virtual std::unique_ptr<ResultType> buildIdHackPlan(const IndexDescriptor* descriptor, + QueryPlannerParams* plannerParams) = 0; + + /** + * Constructs a PlanStage tree from a cached plan and also: + * * Either modifies the constructed tree to run a trial period in order to evaluate the + * cost of a cached plan. If the cost is unexpectedly high, the plan cache entry is + * deactivated and we use multi-planning to select an entirely new winning plan. + * * Or stores additional information in the result object, in case runtime planning is + * implemented as a standalone component, rather than as part of the execution tree. + */ + virtual std::unique_ptr<ResultType> buildCachedPlan(std::unique_ptr<QuerySolution> solution, + const QueryPlannerParams& plannerParams, + size_t decisionWorks) = 0; + + /** + * Constructs a special PlanStage tree for rooted $or queries. Each clause of the $or is planned + * individually, and then an overall query plan is created based on the winning plan from each + * clause. + * + * If sub-planning is implemented as a standalone component, rather than as part of the + * execution tree, this method can populate the result object with additional information + * required to perform the sub-planning. + */ + virtual std::unique_ptr<ResultType> buildSubPlan(const QueryPlannerParams& plannerParams) = 0; + + /** + * If the query have multiple solutions, this method either: + * * Constructs a special PlanStage tree to perform a multi-planning task and pick the best + * plan in runtime. + * * Or builds a PlanStage tree for each of the 'solutions' and stores them in the result + * object, if multi-planning is implemented as a standalone component. + */ + virtual std::unique_ptr<ResultType> buildMultiPlan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + const QueryPlannerParams& plannerParams) = 0; + + OperationContext* _opCtx; + Collection* _collection; + CanonicalQuery* _cq; + PlanYieldPolicy* _yieldPolicy; + const size_t _plannerOptions; +}; + +/** + * A helper class to prepare a classic PlanStage tree for execution. + */ +class ClassicPrepareExecutionHelper final + : public PrepareExecutionHelper<std::unique_ptr<PlanStage>, ClassicPrepareExecutionResult> { +public: + ClassicPrepareExecutionHelper(OperationContext* opCtx, + Collection* collection, + WorkingSet* ws, + CanonicalQuery* cq, + PlanYieldPolicy* yieldPolicy, + size_t plannerOptions) + : PrepareExecutionHelper{opCtx, collection, std::move(cq), yieldPolicy, plannerOptions}, + _ws{ws} {} + +protected: + std::unique_ptr<PlanStage> buildExecutableTree(const QuerySolution& solution) const final { + return stage_builder::buildClassicExecutableTree(_opCtx, _collection, *_cq, solution, _ws); + } + + std::unique_ptr<ClassicPrepareExecutionResult> buildEofPlan() final { + auto result = makeResult(); + result->emplace(std::make_unique<EOFStage>(_cq->getExpCtxRaw()), nullptr); + return result; + } + + std::unique_ptr<ClassicPrepareExecutionResult> buildIdHackPlan( + const IndexDescriptor* descriptor, QueryPlannerParams* plannerParams) final { + auto result = makeResult(); + std::unique_ptr<PlanStage> stage = + std::make_unique<IDHackStage>(_cq->getExpCtxRaw(), _cq, _ws, _collection, descriptor); // Might have to filter out orphaned docs. - if (plannerParams.options & QueryPlannerParams::INCLUDE_SHARD_FILTER) { - root = std::make_unique<ShardFilterStage>( - canonicalQuery->getExpCtxRaw(), - CollectionShardingState::get(opCtx, canonicalQuery->nss()) + if (plannerParams->options & QueryPlannerParams::INCLUDE_SHARD_FILTER) { + stage = std::make_unique<ShardFilterStage>( + _cq->getExpCtxRaw(), + CollectionShardingState::get(_opCtx, _cq->nss()) ->getOwnershipFilter( - opCtx, + _opCtx, CollectionShardingState::OrphanCleanupPolicy::kDisallowOrphanCleanup), - ws, - std::move(root)); + _ws, + std::move(stage)); } - const auto* cqProjection = canonicalQuery->getProj(); + const auto* cqProjection = _cq->getProj(); // Add a SortKeyGeneratorStage if the query requested sortKey metadata. - if (canonicalQuery->metadataDeps()[DocumentMetadataFields::kSortKey]) { - root = std::make_unique<SortKeyGeneratorStage>( - canonicalQuery->getExpCtxRaw(), - std::move(root), - ws, - canonicalQuery->getQueryRequest().getSort()); + if (_cq->metadataDeps()[DocumentMetadataFields::kSortKey]) { + stage = std::make_unique<SortKeyGeneratorStage>( + _cq->getExpCtxRaw(), std::move(stage), _ws, _cq->getQueryRequest().getSort()); } - if (canonicalQuery->getQueryRequest().returnKey()) { - // If returnKey was requested, add ReturnKeyStage to return only the index keys in the - // resulting documents. If a projection was also specified, it will be ignored, with - // the exception the $meta sortKey projection, which can be used along with the + if (_cq->getQueryRequest().returnKey()) { + // If returnKey was requested, add ReturnKeyStage to return only the index keys in + // the resulting documents. If a projection was also specified, it will be ignored, + // with the exception the $meta sortKey projection, which can be used along with the // returnKey. - root = std::make_unique<ReturnKeyStage>( - canonicalQuery->getExpCtxRaw(), + stage = std::make_unique<ReturnKeyStage>( + _cq->getExpCtxRaw(), cqProjection ? QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(*cqProjection) : std::vector<FieldPath>{}, - ws, - std::move(root)); + _ws, + std::move(stage)); } else if (cqProjection) { // There might be a projection. The idhack stage will always fetch the full // document, so we don't support covered projections. However, we might use the // simple inclusion fast path. // Stuff the right data into the params depending on what proj impl we use. if (!cqProjection->isSimple()) { - root = std::make_unique<ProjectionStageDefault>( - canonicalQuery->getExpCtx(), - canonicalQuery->getQueryRequest().getProj(), - canonicalQuery->getProj(), - ws, - std::move(root)); + stage = std::make_unique<ProjectionStageDefault>(_cq->getExpCtxRaw(), + _cq->getQueryRequest().getProj(), + _cq->getProj(), + _ws, + std::move(stage)); } else { - root = std::make_unique<ProjectionStageSimple>( - canonicalQuery->getExpCtxRaw(), - canonicalQuery->getQueryRequest().getProj(), - canonicalQuery->getProj(), - ws, - std::move(root)); + stage = std::make_unique<ProjectionStageSimple>(_cq->getExpCtxRaw(), + _cq->getQueryRequest().getProj(), + _cq->getProj(), + _ws, + std::move(stage)); } } - return PrepareExecutionResult(std::move(canonicalQuery), nullptr, std::move(root)); - } + result->emplace(std::move(stage), nullptr); + return result; + } + + std::unique_ptr<ClassicPrepareExecutionResult> buildCachedPlan( + std::unique_ptr<QuerySolution> solution, + const QueryPlannerParams& plannerParams, + size_t decisionWorks) final { + auto result = makeResult(); + auto&& root = buildExecutableTree(*solution); + + // Add a CachedPlanStage on top of the previous root. + // + // 'decisionWorks' is used to determine whether the existing cache entry should + // be evicted, and the query replanned. + result->emplace(std::make_unique<CachedPlanStage>(_cq->getExpCtxRaw(), + _collection, + _ws, + _cq, + plannerParams, + decisionWorks, + std::move(root)), + std::move(solution)); + return result; + } + + std::unique_ptr<ClassicPrepareExecutionResult> buildSubPlan( + const QueryPlannerParams& plannerParams) final { + auto result = makeResult(); + result->emplace(std::make_unique<SubplanStage>( + _cq->getExpCtxRaw(), _collection, _ws, plannerParams, _cq), + nullptr); + return result; + } + + std::unique_ptr<ClassicPrepareExecutionResult> buildMultiPlan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + const QueryPlannerParams& plannerParams) final { + // 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 multiPlanStage = + std::make_unique<MultiPlanStage>(_cq->getExpCtxRaw(), _collection, _cq); - // Tailable: If the query requests tailable the collection must be capped. - if (canonicalQuery->getQueryRequest().isTailable()) { - if (!collection->isCapped()) { - return Status(ErrorCodes::BadValue, - "error processing query: " + canonicalQuery->toString() + - " tailable cursor requested on non capped collection"); + for (size_t ix = 0; ix < solutions.size(); ++ix) { + if (solutions[ix]->cacheData.get()) { + solutions[ix]->cacheData->indexFilterApplied = plannerParams.indexFiltersApplied; + } + + auto&& nextPlanRoot = buildExecutableTree(*solutions[ix]); + + // Takes ownership of 'nextPlanRoot'. + multiPlanStage->addPlan(std::move(solutions[ix]), std::move(nextPlanRoot), _ws); } + + auto result = makeResult(); + result->emplace(std::move(multiPlanStage), nullptr); + return result; } - // Check that the query should be cached. - if (CollectionQueryInfo::get(collection).getPlanCache()->shouldCacheQuery(*canonicalQuery)) { - // Fill in opDebug information. - const auto planCacheKey = - CollectionQueryInfo::get(collection).getPlanCache()->computeKey(*canonicalQuery); - CurOp::get(opCtx)->debug().queryHash = - canonical_query_encoder::computeHash(planCacheKey.getStableKeyStringData()); - CurOp::get(opCtx)->debug().planCacheKey = - canonical_query_encoder::computeHash(planCacheKey.toString()); - - // Try to look up a cached solution for the query. - if (auto cs = CollectionQueryInfo::get(collection) - .getPlanCache() - ->getCacheEntryIfActive(planCacheKey)) { - // We have a CachedSolution. Have the planner turn it into a QuerySolution. - auto statusWithQs = QueryPlanner::planFromCache(*canonicalQuery, plannerParams, *cs); - - if (statusWithQs.isOK()) { - auto querySolution = std::move(statusWithQs.getValue()); - if ((plannerParams.options & QueryPlannerParams::IS_COUNT) && - turnIxscanIntoCount(querySolution.get())) { - LOGV2_DEBUG(20923, - 2, - "Using fast count: {query}", - "Using fast count", - "query"_attr = redact(canonicalQuery->toStringShort())); - } +private: + WorkingSet* _ws; +}; - auto root = - StageBuilder::build(opCtx, collection, *canonicalQuery, *querySolution, ws); - - // Add a CachedPlanStage on top of the previous root. - // - // 'decisionWorks' is used to determine whether the existing cache entry should - // be evicted, and the query replanned. - auto cachedPlanStage = - std::make_unique<CachedPlanStage>(canonicalQuery->getExpCtxRaw(), - collection, - ws, - canonicalQuery.get(), - plannerParams, - cs->decisionWorks, - std::move(root)); - return PrepareExecutionResult(std::move(canonicalQuery), - std::move(querySolution), - std::move(cachedPlanStage)); +/** + * A helper class to prepare an SBE PlanStage tree for execution. + */ +class SlotBasedPrepareExecutionHelper final + : public PrepareExecutionHelper< + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData>, + SlotBasedPrepareExecutionResult> { +public: + using PrepareExecutionHelper::PrepareExecutionHelper; + +protected: + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> buildExecutableTree( + const QuerySolution& solution) const final { + return buildExecutableTree(solution, false); + } + + std::unique_ptr<SlotBasedPrepareExecutionResult> buildEofPlan() final { + auto result = makeResult(); + result->emplace( + {sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 0, boost::none), + stage_builder::PlanStageData{}}, + nullptr); + return result; + } + + std::unique_ptr<SlotBasedPrepareExecutionResult> buildIdHackPlan( + const IndexDescriptor* descriptor, QueryPlannerParams* plannerParams) final { + uassert(4822862, + "IDHack plan is not supprted by SBE yet", + !(_cq->metadataDeps()[DocumentMetadataFields::kSortKey] || + _cq->getQueryRequest().returnKey() || _cq->getProj())); + + // Fall back to normal planning. + return nullptr; + } + + std::unique_ptr<SlotBasedPrepareExecutionResult> buildCachedPlan( + std::unique_ptr<QuerySolution> solution, + const QueryPlannerParams& plannerParams, + size_t decisionWorks) final { + auto result = makeResult(); + result->emplace(buildExecutableTree(*solution, true), std::move(solution)); + result->setDecisionWorks(decisionWorks); + return result; + } + + std::unique_ptr<SlotBasedPrepareExecutionResult> buildSubPlan( + const QueryPlannerParams& plannerParams) final { + // Nothing do be done here, all planning and stage building will be done by a SubPlanner. + auto result = makeResult(); + result->setNeedsSubplanning(true); + return result; + } + + std::unique_ptr<SlotBasedPrepareExecutionResult> buildMultiPlan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + const QueryPlannerParams& plannerParams) final { + auto result = makeResult(); + for (size_t ix = 0; ix < solutions.size(); ++ix) { + if (solutions[ix]->cacheData.get()) { + solutions[ix]->cacheData->indexFilterApplied = plannerParams.indexFiltersApplied; } + + result->emplace(buildExecutableTree(*solutions[ix], true), std::move(solutions[ix])); } + return result; } - if (internalQueryPlanOrChildrenIndependently.load() && - SubplanStage::canUseSubplanning(*canonicalQuery)) { - LOGV2_DEBUG(20924, - 2, - "Running query as sub-queries: {query}", - "Running query as sub-queries", - "query"_attr = redact(canonicalQuery->toStringShort())); - - root = std::make_unique<SubplanStage>( - canonicalQuery->getExpCtxRaw(), collection, ws, plannerParams, canonicalQuery.get()); - return PrepareExecutionResult(std::move(canonicalQuery), nullptr, std::move(root)); +private: + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> buildExecutableTree( + const QuerySolution& solution, bool needsTrialRunProgressTracker) const { + return stage_builder::buildSlotBasedExecutableTree( + _opCtx, _collection, *_cq, solution, _yieldPolicy, needsTrialRunProgressTracker); } +}; - auto statusWithSolutions = QueryPlanner::plan(*canonicalQuery, plannerParams); - if (!statusWithSolutions.isOK()) { - return statusWithSolutions.getStatus().withContext( - str::stream() << "error processing query: " << canonicalQuery->toString() - << " planner returned error"); +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getClassicExecutor( + OperationContext* opCtx, + Collection* collection, + std::unique_ptr<CanonicalQuery> canonicalQuery, + PlanYieldPolicy::YieldPolicy yieldPolicy, + size_t plannerOptions) { + auto ws = std::make_unique<WorkingSet>(); + ClassicPrepareExecutionHelper helper{ + opCtx, collection, ws.get(), canonicalQuery.get(), nullptr, plannerOptions}; + auto executionResult = helper.prepare(); + if (!executionResult.isOK()) { + return executionResult.getStatus(); } - auto solutions = std::move(statusWithSolutions.getValue()); - // The planner should have returned an error status if there are no solutions. - invariant(solutions.size() > 0); - - // See if one of our solutions is a fast count hack in disguise. - if (plannerParams.options & QueryPlannerParams::IS_COUNT) { - for (size_t i = 0; i < solutions.size(); ++i) { - if (turnIxscanIntoCount(solutions[i].get())) { - // We're not going to cache anything that's fast count. - auto root = - StageBuilder::build(opCtx, collection, *canonicalQuery, *solutions[i], ws); - - LOGV2_DEBUG(20925, - 2, - "Using fast count: {query}, planSummary: {planSummary}", - "Using fast count", - "query"_attr = redact(canonicalQuery->toStringShort()), - "planSummary"_attr = Explain::getPlanSummary(root.get())); + auto&& result = executionResult.getValue(); + auto&& root = result->root(); + invariant(root); + // We must have a tree of stages in order to have a valid plan executor, but the query + // solution may be null. + return PlanExecutor::make(std::move(canonicalQuery), + std::move(ws), + std::move(root), + collection, + yieldPolicy, + {}, + result->solution()); +} - return PrepareExecutionResult( - std::move(canonicalQuery), std::move(solutions[i]), std::move(root)); - } - } +/** + * Checks if the prepared execution plans require further planning in runtime to pick the best + * plan based on the collected execution stats, and returns a 'RuntimePlanner' instance if such + * planning needs to be done, or nullptr otherwise. + */ +std::unique_ptr<sbe::RuntimePlanner> makeRuntimePlannerIfNeeded( + OperationContext* opCtx, + Collection* collection, + CanonicalQuery* canonicalQuery, + size_t numSolutions, + boost::optional<size_t> decisionWorks, + bool needsSubplanning, + PlanYieldPolicySBE* yieldPolicy, + size_t plannerOptions) { + + // If we have multiple solutions, we always need to do the runtime planning. + if (numSolutions > 1) { + invariant(!needsSubplanning && !decisionWorks); + return std::make_unique<sbe::MultiPlanner>( + opCtx, collection, *canonicalQuery, PlanCachingMode::AlwaysCache, yieldPolicy); } - if (1 == solutions.size()) { - // Only one possible plan. Run it. Build the stages from the solution. - auto root = StageBuilder::build(opCtx, collection, *canonicalQuery, *solutions[0], ws); + // If the query can be run as sub-queries, the needSubplanning flag will be set to true and + // we'll need to create a runtime planner to build a composite solution and pick the best plan + // for each sub-query. + if (needsSubplanning) { + invariant(numSolutions == 0); - LOGV2_DEBUG(20926, - 2, - "Only one plan is available; it will be run but will not be cached. {query}, " - "planSummary: {planSummary}", - "Only one plan is available; it will be run but will not be cached", - "query"_attr = redact(canonicalQuery->toStringShort()), - "planSummary"_attr = Explain::getPlanSummary(root.get())); - - return PrepareExecutionResult( - std::move(canonicalQuery), std::move(solutions[0]), std::move(root)); - } else { - // 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 multiPlanStage = std::make_unique<MultiPlanStage>( - canonicalQuery->getExpCtxRaw(), collection, canonicalQuery.get()); + QueryPlannerParams plannerParams; + plannerParams.options = plannerOptions; + fillOutPlannerParams(opCtx, collection, canonicalQuery, &plannerParams); - for (size_t ix = 0; ix < solutions.size(); ++ix) { - if (solutions[ix]->cacheData.get()) { - solutions[ix]->cacheData->indexFilterApplied = plannerParams.indexFiltersApplied; - } + return std::make_unique<sbe::SubPlanner>( + opCtx, collection, *canonicalQuery, plannerParams, yieldPolicy); + } - auto nextPlanRoot = - StageBuilder::build(opCtx, collection, *canonicalQuery, *solutions[ix], ws); + invariant(numSolutions == 1); - // Takes ownership of 'nextPlanRoot'. - multiPlanStage->addPlan(std::move(solutions[ix]), std::move(nextPlanRoot), ws); - } + // If we have a single solution but it was created from a cached plan, we will need to do the + // runtime planning to check if the cached plan still performs efficiently, or requires + // re-planning. The 'decisionWorks' is used to determine whether the existing cache entry should + // be evicted, and the query re-planned. + if (decisionWorks) { + QueryPlannerParams plannerParams; + plannerParams.options = plannerOptions; + fillOutPlannerParams(opCtx, collection, canonicalQuery, &plannerParams); - root = std::move(multiPlanStage); - return PrepareExecutionResult(std::move(canonicalQuery), nullptr, std::move(root)); + return std::make_unique<sbe::CachedSolutionPlanner>( + opCtx, collection, *canonicalQuery, plannerParams, *decisionWorks, yieldPolicy); } -} -} // namespace + // Runtime planning is not required. + return nullptr; +} -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getSlotBasedExecutor( OperationContext* opCtx, Collection* collection, - unique_ptr<CanonicalQuery> canonicalQuery, - PlanExecutor::YieldPolicy yieldPolicy, + std::unique_ptr<CanonicalQuery> cq, + PlanYieldPolicy::YieldPolicy requestedYieldPolicy, size_t plannerOptions) { - unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); - StatusWith<PrepareExecutionResult> executionResult = - prepareExecution(opCtx, collection, ws.get(), std::move(canonicalQuery), plannerOptions); + auto yieldPolicy = + std::make_unique<PlanYieldPolicySBE>(requestedYieldPolicy, + opCtx->getServiceContext()->getFastClockSource(), + internalQueryExecYieldIterations.load(), + Milliseconds{internalQueryExecYieldPeriodMS.load()}); + SlotBasedPrepareExecutionHelper helper{ + opCtx, collection, cq.get(), yieldPolicy.get(), plannerOptions}; + auto executionResult = helper.prepare(); if (!executionResult.isOK()) { return executionResult.getStatus(); } - invariant(executionResult.getValue().root); - // We must have a tree of stages in order to have a valid plan executor, but the query - // solution may be null. - return PlanExecutor::make(std::move(executionResult.getValue().canonicalQuery), - std::move(ws), - std::move(executionResult.getValue().root), - collection, - yieldPolicy, - NamespaceString(), - std::move(executionResult.getValue().querySolution)); + + auto&& result = executionResult.getValue(); + auto&& roots = result->roots(); + auto&& solutions = result->solutions(); + + if (auto planner = makeRuntimePlannerIfNeeded(opCtx, + collection, + cq.get(), + solutions.size(), + result->decisionWorks(), + result->needsSubplanning(), + yieldPolicy.get(), + plannerOptions)) { + // Do the runtime planning and pick the best candidate plan. + auto plan = planner->plan(std::move(solutions), std::move(roots)); + return PlanExecutor::make(opCtx, + std::move(cq), + {std::move(plan.root), std::move(plan.data)}, + {}, + std::move(plan.results), + std::move(yieldPolicy)); + } + // No need for runtime planning, just use the constructed plan stage tree. + invariant(roots.size() == 1); + return PlanExecutor::make( + opCtx, std::move(cq), std::move(roots[0]), {}, std::move(yieldPolicy)); +} +} // namespace + +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor( + OperationContext* opCtx, + Collection* collection, + std::unique_ptr<CanonicalQuery> canonicalQuery, + PlanYieldPolicy::YieldPolicy yieldPolicy, + size_t plannerOptions) { + return internalQueryEnableSlotBasedExecutionEngine.load() + ? getSlotBasedExecutor( + opCtx, collection, std::move(canonicalQuery), yieldPolicy, plannerOptions) + : getClassicExecutor( + opCtx, collection, std::move(canonicalQuery), yieldPolicy, plannerOptions); } // @@ -646,11 +1098,11 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor( namespace { -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> _getExecutorFind( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> _getExecutorFind( OperationContext* opCtx, Collection* collection, - unique_ptr<CanonicalQuery> canonicalQuery, - PlanExecutor::YieldPolicy yieldPolicy, + std::unique_ptr<CanonicalQuery> canonicalQuery, + PlanYieldPolicy::YieldPolicy yieldPolicy, size_t plannerOptions) { if (OperationShardingState::isOperationVersioned(opCtx)) { @@ -661,15 +1113,15 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> _getExecutorFind( } // namespace -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind( OperationContext* opCtx, Collection* collection, - unique_ptr<CanonicalQuery> canonicalQuery, + std::unique_ptr<CanonicalQuery> canonicalQuery, bool permitYield, size_t plannerOptions) { auto yieldPolicy = (permitYield && !opCtx->inMultiDocumentTransaction()) - ? PlanExecutor::YIELD_AUTO - : PlanExecutor::INTERRUPT_ONLY; + ? PlanYieldPolicy::YieldPolicy::YIELD_AUTO + : PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY; return _getExecutorFind( opCtx, collection, std::move(canonicalQuery), yieldPolicy, plannerOptions); } @@ -681,7 +1133,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorLega return _getExecutorFind(opCtx, collection, std::move(canonicalQuery), - PlanExecutor::YIELD_AUTO, + PlanYieldPolicy::YieldPolicy::YIELD_AUTO, QueryPlannerParams::DEFAULT); } @@ -733,7 +1185,7 @@ StatusWith<std::unique_ptr<projection_ast::Projection>> makeProjection(const BSO // Delete // -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete( OpDebug* opDebug, Collection* collection, ParsedDelete* parsedDelete, @@ -771,8 +1223,8 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete( deleteStageParams->opDebug = opDebug; deleteStageParams->stmtId = request->getStmtId(); - unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); - const PlanExecutor::YieldPolicy policy = parsedDelete->yieldPolicy(); + std::unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); + const auto policy = parsedDelete->yieldPolicy(); if (!collection) { // Treat collections that do not exist as empty collections. Return a PlanExecutor which @@ -820,7 +1272,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete( auto idHackStage = std::make_unique<IDHackStage>( expCtx.get(), unparsedQuery["_id"].wrap(), ws.get(), collection, descriptor); - unique_ptr<DeleteStage> root = + std::unique_ptr<DeleteStage> root = std::make_unique<DeleteStage>(expCtx.get(), std::move(deleteStageParams), ws.get(), @@ -840,7 +1292,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete( } // This is the regular path for when we have a CanonicalQuery. - unique_ptr<CanonicalQuery> cq(parsedDelete->releaseParsedQuery()); + std::unique_ptr<CanonicalQuery> cq(parsedDelete->releaseParsedQuery()); // Transfer the explain verbosity level into the expression context. cq->getExpCtx()->explain = verbosity; @@ -861,14 +1313,14 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete( // identify the record to update. const size_t defaultPlannerOptions = QueryPlannerParams::PRESERVE_RECORD_ID; - StatusWith<PrepareExecutionResult> executionResult = - prepareExecution(opCtx, collection, ws.get(), std::move(cq), defaultPlannerOptions); + ClassicPrepareExecutionHelper helper{ + opCtx, collection, ws.get(), cq.get(), nullptr, defaultPlannerOptions}; + auto executionResult = helper.prepare(); if (!executionResult.isOK()) { return executionResult.getStatus(); } - cq = std::move(executionResult.getValue().canonicalQuery); - unique_ptr<QuerySolution> querySolution = std::move(executionResult.getValue().querySolution); - unique_ptr<PlanStage> root = std::move(executionResult.getValue().root); + auto querySolution = executionResult.getValue()->solution(); + auto root = executionResult.getValue()->root(); deleteStageParams->canonicalQuery = cq.get(); @@ -896,7 +1348,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete( // Update // -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate( OpDebug* opDebug, Collection* collection, ParsedUpdate* parsedUpdate, @@ -940,9 +1392,9 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate( str::stream() << "Not primary while performing update on " << nss.ns()); } - const PlanExecutor::YieldPolicy policy = parsedUpdate->yieldPolicy(); + const auto policy = parsedUpdate->yieldPolicy(); - unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); + std::unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); UpdateStageParams updateStageParams(request, driver, opDebug); // If the collection doesn't exist, then return a PlanExecutor for a no-op EOF plan. We have @@ -1006,7 +1458,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate( } // This is the regular path for when we have a CanonicalQuery. - unique_ptr<CanonicalQuery> cq(parsedUpdate->releaseParsedQuery()); + std::unique_ptr<CanonicalQuery> cq(parsedUpdate->releaseParsedQuery()); std::unique_ptr<projection_ast::Projection> projection; if (!request->getProj().isEmpty()) { @@ -1027,14 +1479,14 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate( // identify the record to update. const size_t defaultPlannerOptions = QueryPlannerParams::PRESERVE_RECORD_ID; - StatusWith<PrepareExecutionResult> executionResult = - prepareExecution(opCtx, collection, ws.get(), std::move(cq), defaultPlannerOptions); + ClassicPrepareExecutionHelper helper{ + opCtx, collection, ws.get(), cq.get(), nullptr, defaultPlannerOptions}; + auto executionResult = helper.prepare(); if (!executionResult.isOK()) { return executionResult.getStatus(); } - cq = std::move(executionResult.getValue().canonicalQuery); - unique_ptr<QuerySolution> querySolution = std::move(executionResult.getValue().querySolution); - unique_ptr<PlanStage> root = std::move(executionResult.getValue().root); + auto querySolution = executionResult.getValue()->solution(); + auto root = executionResult.getValue()->root(); invariant(root); updateStageParams.canonicalQuery = cq.get(); @@ -1199,14 +1651,14 @@ bool getDistinctNodeIndex(const std::vector<IndexEntry>& indices, } // namespace -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount( const boost::intrusive_ptr<ExpressionContext>& expCtx, Collection* collection, const CountCommand& request, bool explain, const NamespaceString& nss) { OperationContext* opCtx = expCtx->opCtx; - unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); + std::unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); auto qr = std::make_unique<QueryRequest>(nss); qr->setFilter(request.getQuery()); @@ -1227,10 +1679,11 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount( if (!statusWithCQ.isOK()) { return statusWithCQ.getStatus(); } - unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); + std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); - const auto yieldPolicy = opCtx->inMultiDocumentTransaction() ? PlanExecutor::INTERRUPT_ONLY - : PlanExecutor::YIELD_AUTO; + const auto yieldPolicy = opCtx->inMultiDocumentTransaction() + ? PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY + : PlanYieldPolicy::YieldPolicy::YIELD_AUTO; const auto skip = request.getSkip().value_or(0); const auto limit = request.getLimit().value_or(0); @@ -1239,7 +1692,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount( // Treat collections that do not exist as empty collections. Note that the explain reporting // machinery always assumes that the root stage for a count operation is a CountStage, so in // this case we put a CountStage on top of an EOFStage. - unique_ptr<PlanStage> root = std::make_unique<CountStage>( + std::unique_ptr<PlanStage> root = std::make_unique<CountStage>( expCtx.get(), collection, limit, skip, ws.get(), new EOFStage(expCtx.get())); return PlanExecutor::make( expCtx, std::move(ws), std::move(root), nullptr, yieldPolicy, nss); @@ -1255,7 +1708,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount( const bool useRecordStoreCount = isEmptyQueryPredicate && request.getHint().isEmpty(); if (useRecordStoreCount) { - unique_ptr<PlanStage> root = + std::unique_ptr<PlanStage> root = std::make_unique<RecordStoreFastCountStage>(expCtx.get(), collection, skip, limit); return PlanExecutor::make( expCtx, std::move(ws), std::move(root), nullptr, yieldPolicy, nss); @@ -1266,14 +1719,14 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount( plannerOptions |= QueryPlannerParams::INCLUDE_SHARD_FILTER; } - StatusWith<PrepareExecutionResult> executionResult = - prepareExecution(opCtx, collection, ws.get(), std::move(cq), plannerOptions); + ClassicPrepareExecutionHelper helper{ + opCtx, collection, ws.get(), cq.get(), nullptr, plannerOptions}; + auto executionResult = helper.prepare(); if (!executionResult.isOK()) { return executionResult.getStatus(); } - cq = std::move(executionResult.getValue().canonicalQuery); - unique_ptr<QuerySolution> querySolution = std::move(executionResult.getValue().querySolution); - unique_ptr<PlanStage> root = std::move(executionResult.getValue().root); + auto querySolution = executionResult.getValue()->solution(); + auto root = executionResult.getValue()->root(); invariant(root); @@ -1296,7 +1749,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount( // bool turnIxscanIntoDistinctIxscan(QuerySolution* soln, - const string& field, + const std::string& field, bool strictDistinctOnly) { QuerySolutionNode* root = soln->root.get(); @@ -1553,7 +2006,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorForS OperationContext* opCtx, Collection* collection, const QueryPlannerParams& plannerParams, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, ParsedDistinct* parsedDistinct) { invariant(parsedDistinct->getQuery()); auto collator = parsedDistinct->getQuery()->getCollator(); @@ -1591,9 +2044,9 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorForS *parsedDistinct->getQuery(), params, std::move(solnRoot)); invariant(soln); - unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); - auto root = - StageBuilder::build(opCtx, collection, *parsedDistinct->getQuery(), *soln, ws.get()); + std::unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); + auto&& root = stage_builder::buildClassicExecutableTree( + opCtx, collection, *parsedDistinct->getQuery(), *soln, ws.get()); LOGV2_DEBUG(20931, 2, @@ -1611,7 +2064,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorForS std::move(soln)); } -// Checks each solution in the 'solutions' vector to see if one includes an IXSCAN that can be +// Checks each solution in the 'solutions' std::vector to see if one includes an IXSCAN that can be // rewritten as a DISTINCT_SCAN, assuming we want distinct scan behavior on the getKey() property of // the 'parsedDistinct' argument. // @@ -1627,7 +2080,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDistinctFromIndexSolutions(OperationContext* opCtx, Collection* collection, std::vector<std::unique_ptr<QuerySolution>> solutions, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, ParsedDistinct* parsedDistinct, bool strictDistinctOnly) { // We look for a solution that has an ixscan we can turn into a distinctixscan @@ -1635,9 +2088,9 @@ getExecutorDistinctFromIndexSolutions(OperationContext* opCtx, if (turnIxscanIntoDistinctIxscan( solutions[i].get(), parsedDistinct->getKey(), strictDistinctOnly)) { // Build and return the SSR over solutions[i]. - unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); - unique_ptr<QuerySolution> currentSolution = std::move(solutions[i]); - auto root = StageBuilder::build( + std::unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); + std::unique_ptr<QuerySolution> currentSolution = std::move(solutions[i]); + auto&& root = stage_builder::buildClassicExecutableTree( opCtx, collection, *parsedDistinct->getQuery(), *currentSolution, ws.get()); LOGV2_DEBUG(20932, @@ -1664,11 +2117,11 @@ getExecutorDistinctFromIndexSolutions(OperationContext* opCtx, /** * Makes a clone of 'cq' but without any projection, then runs getExecutor on the clone. */ -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorWithoutProjection( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorWithoutProjection( OperationContext* opCtx, Collection* collection, const CanonicalQuery* cq, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, size_t plannerOptions) { auto qr = std::make_unique<QueryRequest>(cq->getQueryRequest()); qr->setProj(BSONObj()); @@ -1687,12 +2140,13 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorWithoutPr } } // namespace -StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDistinct( +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDistinct( Collection* collection, size_t plannerOptions, ParsedDistinct* parsedDistinct) { auto expCtx = parsedDistinct->getQuery()->getExpCtx(); OperationContext* opCtx = expCtx->opCtx; - const auto yieldPolicy = opCtx->inMultiDocumentTransaction() ? PlanExecutor::INTERRUPT_ONLY - : PlanExecutor::YIELD_AUTO; + const auto yieldPolicy = opCtx->inMultiDocumentTransaction() + ? PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY + : PlanYieldPolicy::YieldPolicy::YIELD_AUTO; if (!collection) { // Treat collections that do not exist as empty collections. diff --git a/src/mongo/db/query/get_executor.h b/src/mongo/db/query/get_executor.h index 905903bd607..7f39d4436c2 100644 --- a/src/mongo/db/query/get_executor.h +++ b/src/mongo/db/query/get_executor.h @@ -121,7 +121,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor( OperationContext* opCtx, Collection* collection, std::unique_ptr<CanonicalQuery> canonicalQuery, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, size_t plannerOptions = 0); /** diff --git a/src/mongo/db/query/index_bounds_builder.cpp b/src/mongo/db/query/index_bounds_builder.cpp index 1f8dcae97e9..faed197b1c6 100644 --- a/src/mongo/db/query/index_bounds_builder.cpp +++ b/src/mongo/db/query/index_bounds_builder.cpp @@ -1212,6 +1212,54 @@ void IndexBoundsBuilder::alignBounds(IndexBounds* bounds, const BSONObj& kp, int } } +void IndexBoundsBuilder::appendTrailingAllValuesInterval(const Interval& interval, + bool startKeyInclusive, + bool endKeyInclusive, + BSONObjBuilder* startBob, + BSONObjBuilder* endBob) { + invariant(startBob); + invariant(endBob); + + // Must be min->max or max->min. + if (interval.isMinToMax()) { + // As an example for the logic below, consider the index {a:1, b:1} and a count for + // {a: {$gt: 2}}. Our start key isn't inclusive (as it's $gt: 2) and looks like + // {"":2} so far. If we move to the key greater than {"":2, "": MaxKey} we will get + // the first value of 'a' that is greater than 2. + if (!startKeyInclusive) { + startBob->appendMaxKey(""); + } else { + // In this case, consider the index {a:1, b:1} and a count for {a:{$gte: 2}}. + // We want to look at all values where a is 2, so our start key is {"":2, + // "":MinKey}. + startBob->appendMinKey(""); + } + + // Same deal as above. Consider the index {a:1, b:1} and a count for {a: {$lt: 2}}. + // Our end key isn't inclusive as ($lt: 2) and looks like {"":2} so far. We can't + // look at any values where a is 2 so we have to stop at {"":2, "": MinKey} as + // that's the smallest key where a is still 2. + if (!endKeyInclusive) { + endBob->appendMinKey(""); + } else { + endBob->appendMaxKey(""); + } + } else if (interval.isMaxToMin()) { + // The reasoning here is the same as above but with the directions reversed. + if (!startKeyInclusive) { + startBob->appendMinKey(""); + } else { + startBob->appendMaxKey(""); + } + + if (!endKeyInclusive) { + endBob->appendMaxKey(""); + } else { + endBob->appendMinKey(""); + } + } +} + // static bool IndexBoundsBuilder::isSingleInterval(const IndexBounds& bounds, BSONObj* startKey, @@ -1265,12 +1313,6 @@ bool IndexBoundsBuilder::isSingleInterval(const IndexBounds& bounds, ++fieldNo; - // Get some "all values" intervals for comparison's sake. - // TODO: make static? - Interval minMax = IndexBoundsBuilder::allValues(); - Interval maxMin = minMax; - maxMin.reverse(); - // And after the non-point interval we can have any number of "all values" intervals. for (; fieldNo < bounds.fields.size(); ++fieldNo) { const OrderedIntervalList& oil = bounds.fields[fieldNo]; @@ -1279,42 +1321,9 @@ bool IndexBoundsBuilder::isSingleInterval(const IndexBounds& bounds, break; } - // Must be min->max or max->min. - if (oil.intervals[0].equals(minMax)) { - // As an example for the logic below, consider the index {a:1, b:1} and a count for - // {a: {$gt: 2}}. Our start key isn't inclusive (as it's $gt: 2) and looks like - // {"":2} so far. If we move to the key greater than {"":2, "": MaxKey} we will get - // the first value of 'a' that is greater than 2. - if (!*startKeyInclusive) { - startBob.appendMaxKey(""); - } else { - // In this case, consider the index {a:1, b:1} and a count for {a:{$gte: 2}}. - // We want to look at all values where a is 2, so our start key is {"":2, - // "":MinKey}. - startBob.appendMinKey(""); - } - - // Same deal as above. Consider the index {a:1, b:1} and a count for {a: {$lt: 2}}. - // Our end key isn't inclusive as ($lt: 2) and looks like {"":2} so far. We can't - // look at any values where a is 2 so we have to stop at {"":2, "": MinKey} as - // that's the smallest key where a is still 2. - if (!*endKeyInclusive) { - endBob.appendMinKey(""); - } else { - endBob.appendMaxKey(""); - } - } else if (oil.intervals[0].equals(maxMin)) { - // The reasoning here is the same as above but with the directions reversed. - if (!*startKeyInclusive) { - startBob.appendMinKey(""); - } else { - startBob.appendMaxKey(""); - } - if (!*endKeyInclusive) { - endBob.appendMaxKey(""); - } else { - endBob.appendMinKey(""); - } + if (oil.intervals[0].isMinToMax() || oil.intervals[0].isMaxToMin()) { + IndexBoundsBuilder::appendTrailingAllValuesInterval( + oil.intervals[0], *startKeyInclusive, *endKeyInclusive, &startBob, &endBob); } else { // No dice. break; diff --git a/src/mongo/db/query/index_bounds_builder.h b/src/mongo/db/query/index_bounds_builder.h index 0453df842cb..7f141600837 100644 --- a/src/mongo/db/query/index_bounds_builder.h +++ b/src/mongo/db/query/index_bounds_builder.h @@ -208,6 +208,36 @@ public: BSONObj* endKey, bool* endKeyInclusive); + /** + * Appends the startKey and endKey of the given "all values" 'interval' (which is either + * [MinKey, MaxKey] or [MaxKey, MinKey] interval) to the 'startBob' and 'endBob' respectively, + * handling inclusivity of each bound through the relevant '*KeyInclusive' parameter. + * + * If the 'interval' is not an "all values" interval, does nothing. + * + * Precondition: startBob and endBob should contain one or more leading intervals which are not + * "all values" intervals, to make the constructed interval valid. + * + * The decision whether to append MinKey or MaxKey value either to startBob or endBob is based + * on the interval type (min -> max or max -> min), and inclusivity flags. + * + * As an example, consider the index {a:1, b:1} and a count for {a: {$gt: 2}}. Our start key + * isn't inclusive (as it's $gt: 2) and looks like {"":2} so far. Because {a: 2, b: MaxKey} + * sorts *after* any real-world data pair {a: 2, b: anything}, setting it as the start value + * ensures that the first index entry we encounter will be the smallest key with a > 2. + * + * Same logic applies if the end key is not inclusive. Consider the index {a:1, b:1} and a count + * for {a: {$lt: 2}}. Our end key isn't inclusive as ($lt: 2) and looks like {"":2} so far. + * Because {a: 2, b: MinKey} sorts *before* any real-world data pair {a: 2, b: anything}, + * setting it as the end value ensures that the final index entry we encounter will be the last + * key with a < 2. + */ + static void appendTrailingAllValuesInterval(const Interval& interval, + bool startKeyInclusive, + bool endKeyInclusive, + BSONObjBuilder* startBob, + BSONObjBuilder* endBob); + private: /** * Performs the heavy lifting for IndexBoundsBuilder::translate(). diff --git a/src/mongo/db/query/internal_plans.cpp b/src/mongo/db/query/internal_plans.cpp index b826311614f..33d2a3d1489 100644 --- a/src/mongo/db/query/internal_plans.cpp +++ b/src/mongo/db/query/internal_plans.cpp @@ -51,7 +51,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::collection OperationContext* opCtx, StringData ns, Collection* collection, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, const Direction direction) { std::unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); @@ -82,7 +82,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWith OperationContext* opCtx, Collection* collection, std::unique_ptr<DeleteStageParams> params, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, Direction direction) { invariant(collection); auto ws = std::make_unique<WorkingSet>(); @@ -109,7 +109,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::indexScan( const BSONObj& startKey, const BSONObj& endKey, BoundInclusion boundInclusion, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, Direction direction, int options) { auto ws = std::make_unique<WorkingSet>(); @@ -141,7 +141,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWith const BSONObj& startKey, const BSONObj& endKey, BoundInclusion boundInclusion, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, Direction direction) { invariant(collection); auto ws = std::make_unique<WorkingSet>(); @@ -174,7 +174,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::updateWith const UpdateStageParams& params, const IndexDescriptor* descriptor, const BSONObj& key, - PlanExecutor::YieldPolicy yieldPolicy) { + PlanYieldPolicy::YieldPolicy yieldPolicy) { invariant(collection); auto ws = std::make_unique<WorkingSet>(); diff --git a/src/mongo/db/query/internal_plans.h b/src/mongo/db/query/internal_plans.h index 228f6f4fd97..05e357fadd1 100644 --- a/src/mongo/db/query/internal_plans.h +++ b/src/mongo/db/query/internal_plans.h @@ -72,7 +72,7 @@ public: OperationContext* opCtx, StringData ns, Collection* collection, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, const Direction direction = FORWARD); /** @@ -82,7 +82,7 @@ public: OperationContext* opCtx, Collection* collection, std::unique_ptr<DeleteStageParams> params, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, Direction direction = FORWARD); /** @@ -95,7 +95,7 @@ public: const BSONObj& startKey, const BSONObj& endKey, BoundInclusion boundInclusion, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, Direction direction = FORWARD, int options = IXSCAN_DEFAULT); @@ -110,7 +110,7 @@ public: const BSONObj& startKey, const BSONObj& endKey, BoundInclusion boundInclusion, - PlanExecutor::YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, Direction direction = FORWARD); /** @@ -122,7 +122,7 @@ public: const UpdateStageParams& params, const IndexDescriptor* descriptor, const BSONObj& key, - PlanExecutor::YieldPolicy yieldPolicy); + PlanYieldPolicy::YieldPolicy yieldPolicy); private: /** diff --git a/src/mongo/db/query/interval.cpp b/src/mongo/db/query/interval.cpp index 1aa1e816bbe..7462d5ee01d 100644 --- a/src/mongo/db/query/interval.cpp +++ b/src/mongo/db/query/interval.cpp @@ -168,6 +168,10 @@ bool Interval::isMinToMax() const { return (start.type() == BSONType::MinKey && end.type() == BSONType::MaxKey); } +bool Interval::isMaxToMin() const { + return (start.type() == BSONType::MaxKey && end.type() == BSONType::MinKey); +} + Interval::IntervalComparison Interval::compare(const Interval& other) const { // // Intersect cases diff --git a/src/mongo/db/query/interval.h b/src/mongo/db/query/interval.h index e86af179f78..7a9c565d3f0 100644 --- a/src/mongo/db/query/interval.h +++ b/src/mongo/db/query/interval.h @@ -140,6 +140,11 @@ struct Interval { */ bool isMinToMax() const; + /** + * Returns true if the interval is from MaxKey to MinKey. + */ + bool isMaxToMin() const; + /** Returns how 'this' compares to 'other' */ enum IntervalComparison { // diff --git a/src/mongo/db/query/mock_yield_policies.h b/src/mongo/db/query/mock_yield_policies.h index 68d537f34a3..69661392b46 100644 --- a/src/mongo/db/query/mock_yield_policies.h +++ b/src/mongo/db/query/mock_yield_policies.h @@ -30,7 +30,7 @@ #pragma once #include "mongo/base/error_codes.h" -#include "mongo/db/query/plan_yield_policy.h" +#include "mongo/db/query/plan_executor.h" namespace mongo { @@ -38,19 +38,22 @@ namespace mongo { * A custom yield policy that always reports the plan should yield, and always returns * ErrorCodes::ExceededTimeLimit from yield(). */ -class AlwaysTimeOutYieldPolicy : public PlanYieldPolicy { +class AlwaysTimeOutYieldPolicy final : public PlanYieldPolicy { public: AlwaysTimeOutYieldPolicy(PlanExecutor* exec) - : PlanYieldPolicy(exec, PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT) {} + : PlanYieldPolicy(PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT, + exec->getOpCtx()->getServiceContext()->getFastClockSource(), + 0, + Milliseconds{0}) {} AlwaysTimeOutYieldPolicy(ClockSource* cs) - : PlanYieldPolicy(PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT, cs) {} + : PlanYieldPolicy(PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT, cs, 0, Milliseconds{0}) {} - bool shouldYieldOrInterrupt() override { + bool shouldYieldOrInterrupt(OperationContext*) override { return true; } - Status yieldOrInterrupt(std::function<void()> whileYieldingFn) override { + Status yield(OperationContext*, std::function<void()> whileYieldingFn) override { return {ErrorCodes::ExceededTimeLimit, "Using AlwaysTimeOutYieldPolicy"}; } }; @@ -59,21 +62,40 @@ public: * A custom yield policy that always reports the plan should yield, and always returns * ErrorCodes::QueryPlanKilled from yield(). */ -class AlwaysPlanKilledYieldPolicy : public PlanYieldPolicy { +class AlwaysPlanKilledYieldPolicy final : public PlanYieldPolicy { public: AlwaysPlanKilledYieldPolicy(PlanExecutor* exec) - : PlanYieldPolicy(exec, PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED) {} + : PlanYieldPolicy(PlanYieldPolicy::YieldPolicy::ALWAYS_MARK_KILLED, + exec->getOpCtx()->getServiceContext()->getFastClockSource(), + 0, + Milliseconds{0}) {} AlwaysPlanKilledYieldPolicy(ClockSource* cs) - : PlanYieldPolicy(PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED, cs) {} + : PlanYieldPolicy( + PlanYieldPolicy::YieldPolicy::ALWAYS_MARK_KILLED, cs, 0, Milliseconds{0}) {} - bool shouldYieldOrInterrupt() override { + bool shouldYieldOrInterrupt(OperationContext*) override { return true; } - Status yieldOrInterrupt(std::function<void()> whileYieldingFn) override { + Status yield(OperationContext*, std::function<void()> whileYieldingFn) override { return {ErrorCodes::QueryPlanKilled, "Using AlwaysPlanKilledYieldPolicy"}; } }; +class NoopYieldPolicy final : public PlanYieldPolicy { +public: + NoopYieldPolicy(ClockSource* clockSource) + : PlanYieldPolicy(PlanYieldPolicy::YieldPolicy::NO_YIELD, clockSource, 0, Milliseconds{0}) { + } + + bool shouldYieldOrInterrupt(OperationContext*) override { + return false; + } + + Status yield(OperationContext*, std::function<void()> whileYieldingFn) override { + MONGO_UNREACHABLE; + } +}; + } // namespace mongo diff --git a/src/mongo/db/query/plan_cache.cpp b/src/mongo/db/query/plan_cache.cpp index 30c13ce0f5e..03dc254f272 100644 --- a/src/mongo/db/query/plan_cache.cpp +++ b/src/mongo/db/query/plan_cache.cpp @@ -55,6 +55,7 @@ #include "mongo/util/assert_util.h" #include "mongo/util/hex.h" #include "mongo/util/transitional_tools_do_not_use/vector_spooling.h" +#include "mongo/util/visit_helper.h" namespace mongo { namespace { @@ -198,7 +199,7 @@ CachedSolution::~CachedSolution() { std::unique_ptr<PlanCacheEntry> PlanCacheEntry::create( const std::vector<QuerySolution*>& solutions, - std::unique_ptr<const PlanRankingDecision> decision, + std::unique_ptr<const plan_ranker::PlanRankingDecision> decision, const CanonicalQuery& query, uint32_t queryHash, uint32_t planCacheKey, @@ -252,7 +253,7 @@ PlanCacheEntry::PlanCacheEntry(std::vector<std::unique_ptr<const SolutionCacheDa const Date_t timeOfCreation, const uint32_t queryHash, const uint32_t planCacheKey, - std::unique_ptr<const PlanRankingDecision> decision, + std::unique_ptr<const plan_ranker::PlanRankingDecision> decision, const bool isActive, const size_t works) : plannerData(std::move(plannerData)), @@ -283,7 +284,7 @@ std::unique_ptr<PlanCacheEntry> PlanCacheEntry::clone() const { solutionCacheData[i] = std::unique_ptr<const SolutionCacheData>(plannerData[i]->clone()); } - auto decisionPtr = std::unique_ptr<PlanRankingDecision>(decision->clone()); + auto decisionPtr = std::unique_ptr<plan_ranker::PlanRankingDecision>(decision->clone()); return std::unique_ptr<PlanCacheEntry>(new PlanCacheEntry(std::move(solutionCacheData), query, sort, @@ -556,7 +557,7 @@ PlanCache::NewEntryState PlanCache::getNewEntryState(const CanonicalQuery& query Status PlanCache::set(const CanonicalQuery& query, const std::vector<QuerySolution*>& solns, - std::unique_ptr<PlanRankingDecision> why, + std::unique_ptr<plan_ranker::PlanRankingDecision> why, Date_t now, boost::optional<double> worksGrowthCoefficient) { invariant(why); @@ -565,7 +566,8 @@ Status PlanCache::set(const CanonicalQuery& query, return Status(ErrorCodes::BadValue, "no solutions provided"); } - if (why->stats.size() != solns.size()) { + auto statsSize = stdx::visit([](auto&& stats) { return stats.size(); }, why->stats); + if (statsSize != solns.size()) { return Status(ErrorCodes::BadValue, "number of stats in decision must match solutions"); } @@ -580,8 +582,16 @@ Status PlanCache::set(const CanonicalQuery& query, "match the number of solutions"); } + const size_t newWorks = stdx::visit( + visit_helper::Overloaded{[](std::vector<std::unique_ptr<PlanStageStats>>& stats) { + return stats[0]->common.works; + }, + [](std::vector<std::unique_ptr<sbe::PlanStageStats>>& stats) { + return calculateNumberOfReads(stats[0].get()); + }}, + + why->stats); const auto key = computeKey(query); - const size_t newWorks = why->stats[0]->common.works; stdx::lock_guard<Latch> cacheLock(_cacheMutex); bool isNewEntryActive = false; uint32_t queryHash; diff --git a/src/mongo/db/query/plan_cache.h b/src/mongo/db/query/plan_cache.h index c5d3133302d..464370cac35 100644 --- a/src/mongo/db/query/plan_cache.h +++ b/src/mongo/db/query/plan_cache.h @@ -43,7 +43,6 @@ #include "mongo/util/container_size_helper.h" namespace mongo { - /** * Represents the "key" used in the PlanCache mapping from query shape -> query plan. */ @@ -107,8 +106,9 @@ public: } }; - +namespace plan_ranker { struct PlanRankingDecision; +} struct QuerySolution; struct QuerySolutionNode; @@ -313,7 +313,7 @@ public: */ static std::unique_ptr<PlanCacheEntry> create( const std::vector<QuerySolution*>& solutions, - std::unique_ptr<const PlanRankingDecision> decision, + std::unique_ptr<const plan_ranker::PlanRankingDecision> decision, const CanonicalQuery& query, uint32_t queryHash, uint32_t planCacheKey, @@ -363,7 +363,7 @@ public: // // Information that went into picking the winning plan and also why the other plans lost. - const std::unique_ptr<const PlanRankingDecision> decision; + const std::unique_ptr<const plan_ranker::PlanRankingDecision> decision; // Whether or not the cache entry is active. Inactive cache entries should not be used for // planning. @@ -392,7 +392,7 @@ private: Date_t timeOfCreation, uint32_t queryHash, uint32_t planCacheKey, - std::unique_ptr<const PlanRankingDecision> decision, + std::unique_ptr<const plan_ranker::PlanRankingDecision> decision, bool isActive, size_t works); @@ -477,7 +477,7 @@ public: */ Status set(const CanonicalQuery& query, const std::vector<QuerySolution*>& solns, - std::unique_ptr<PlanRankingDecision> why, + std::unique_ptr<plan_ranker::PlanRankingDecision> why, Date_t now, boost::optional<double> worksGrowthCoefficient = boost::none); @@ -594,5 +594,4 @@ private: // are allowed. PlanCacheIndexabilityState _indexabilityState; }; - } // namespace mongo diff --git a/src/mongo/db/query/plan_cache_test.cpp b/src/mongo/db/query/plan_cache_test.cpp index 764dd716978..41337bd9cc5 100644 --- a/src/mongo/db/query/plan_cache_test.cpp +++ b/src/mongo/db/query/plan_cache_test.cpp @@ -296,17 +296,20 @@ struct GenerateQuerySolution { /** * Utility function to create a PlanRankingDecision */ -std::unique_ptr<PlanRankingDecision> createDecision(size_t numPlans, size_t works = 0) { - unique_ptr<PlanRankingDecision> why(new PlanRankingDecision()); +std::unique_ptr<plan_ranker::PlanRankingDecision> createDecision(size_t numPlans, + size_t works = 0) { + auto why = std::make_unique<plan_ranker::PlanRankingDecision>(); + std::vector<std::unique_ptr<PlanStageStats>> stats; for (size_t i = 0; i < numPlans; ++i) { CommonStats common("COLLSCAN"); - auto stats = std::make_unique<PlanStageStats>(common, STAGE_COLLSCAN); - stats->specific.reset(new CollectionScanStats()); - why->stats.push_back(std::move(stats)); - why->stats[i]->common.works = works; + auto stat = std::make_unique<PlanStageStats>(common, STAGE_COLLSCAN); + stat->specific.reset(new CollectionScanStats()); + stat->common.works = works; + stats.push_back(std::move(stat)); why->scores.push_back(0U); why->candidateOrder.push_back(i); } + why->getStats<PlanStageStats>() = std::move(stats); return why; } @@ -492,7 +495,7 @@ TEST(PlanCacheTest, AddEmptySolutions) { PlanCache planCache; unique_ptr<CanonicalQuery> cq(canonicalize("{a: 1}")); std::vector<QuerySolution*> solns; - unique_ptr<PlanRankingDecision> decision(createDecision(1U)); + unique_ptr<plan_ranker::PlanRankingDecision> decision(createDecision(1U)); QueryTestServiceContext serviceContext; ASSERT_NOT_OK(planCache.set(*cq, solns, std::move(decision), Date_t{})); } @@ -678,7 +681,7 @@ TEST(PlanCacheTest, WorksValueIncreases) { ASSERT_EQ(planCache.get(*cq).state, PlanCache::CacheEntryState::kPresentActive); entry = assertGet(planCache.getEntry(*cq)); ASSERT_TRUE(entry->isActive); - ASSERT_EQ(entry->decision->stats[0]->common.works, 25U); + ASSERT_EQ(entry->decision->getStats<PlanStageStats>()[0]->common.works, 25U); ASSERT_EQ(entry->works, 25U); ASSERT_EQUALS(planCache.size(), 1U); diff --git a/src/mongo/db/query/plan_executor.h b/src/mongo/db/query/plan_executor.h index 7c241a03b83..40a8e4fee18 100644 --- a/src/mongo/db/query/plan_executor.h +++ b/src/mongo/db/query/plan_executor.h @@ -30,10 +30,15 @@ #pragma once #include <boost/optional.hpp> +#include <queue> #include "mongo/base/status.h" #include "mongo/db/catalog/util/partitioned.h" +#include "mongo/db/exec/plan_stats.h" +#include "mongo/db/query/plan_yield_policy.h" +#include "mongo/db/query/plan_yield_policy_sbe.h" #include "mongo/db/query/query_solution.h" +#include "mongo/db/query/sbe_stage_builder.h" #include "mongo/db/storage/snapshot.h" #include "mongo/stdx/unordered_set.h" @@ -45,11 +50,13 @@ struct CappedInsertNotifierData; class Collection; class PlanExecutor; class PlanStage; -class PlanYieldPolicy; class RecordId; -struct PlanStageStats; class WorkingSet; +namespace sbe { +class PlanStage; +} // namespace sbe + /** * If a getMore command specified a lastKnownCommittedOpTime (as secondaries do), we want to stop * waiting for new data as soon as the committed op time changes. @@ -78,101 +85,24 @@ public: IS_EOF, }; - /** - * The yielding policy of the plan executor. By default, an executor does not yield itself - * (NO_YIELD). - */ - enum YieldPolicy { - // Any call to getNext() may yield. In particular, the executor may die on any call to - // getNext() due to a required index or collection becoming invalid during yield. If this - // occurs, getNext() will produce an error during yield recovery and will throw an - // exception. Additionally, this will handle all WriteConflictExceptions that occur while - // processing the query. With this yield policy, it is possible for getNext() to return - // throw with locks released. Cleanup that happens while the stack unwinds cannot assume - // locks are held. - YIELD_AUTO, - - // This will handle WriteConflictExceptions that occur while processing the query, but will - // not yield locks. abandonSnapshot() will be called if a WriteConflictException occurs so - // callers must be prepared to get a new snapshot. The caller must hold their locks - // continuously from construction to destruction. Callers which do not want auto-yielding, - // but may release their locks during query execution must use the YIELD_MANUAL policy. - WRITE_CONFLICT_RETRY_ONLY, - - // Use this policy if you want to disable auto-yielding, but will release locks while using - // the PlanExecutor. Any WriteConflictExceptions will be raised to the caller of getNext(). - // - // With this policy, an explicit call must be made to saveState() before releasing locks, - // and an explicit call to restoreState() must be made after reacquiring locks. - // restoreState() will throw if the PlanExecutor is now invalid due to a catalog operation - // (e.g. collection drop) during yield. - YIELD_MANUAL, - - // Can be used in one of the following scenarios: - // - The caller will hold a lock continuously for the lifetime of this PlanExecutor. - // - This PlanExecutor doesn't logically belong to a Collection, and so does not need to be - // locked during execution. For example, a PlanExecutor containing a PipelineProxyStage - // which is being used to execute an aggregation pipeline. - NO_YIELD, - - // Will not yield locks or storage engine resources, but will check for interrupt. - INTERRUPT_ONLY, - - // Used for testing, this yield policy will cause the PlanExecutor to time out on the first - // yield, throwing an ErrorCodes::ExceededTimeLimit error. - ALWAYS_TIME_OUT, - - // Used for testing, this yield policy will cause the PlanExecutor to be marked as killed on - // the first yield, throwing an ErrorCodes::QueryPlanKilled error. - ALWAYS_MARK_KILLED, + // Describes whether callers should acquire locks when using a PlanExecutor. Not all cursors + // have the same locking behavior. In particular, find executors using the legacy PlanStage + // engine require the caller to lock the collection in MODE_IS. Aggregate executors and SBE + // executors, on the other hand, may access multiple collections and acquire their own locks on + // any involved collections while producing query results. Therefore, the caller need not + // explicitly acquire any locks for such PlanExecutors. + // + // The policy is consulted on getMore in order to determine locking behavior, since during + // getMore we otherwise could not easily know what flavor of cursor we're using. + enum class LockPolicy { + // The caller is responsible for locking the collection over which this PlanExecutor + // executes. + kLockExternally, + + // The caller need not hold no locks; this PlanExecutor acquires any necessary locks itself. + kLocksInternally, }; - static std::string serializeYieldPolicy(YieldPolicy yieldPolicy) { - switch (yieldPolicy) { - case YIELD_AUTO: - return "YIELD_AUTO"; - case WRITE_CONFLICT_RETRY_ONLY: - return "WRITE_CONFLICT_RETRY_ONLY"; - case YIELD_MANUAL: - return "YIELD_MANUAL"; - case NO_YIELD: - return "NO_YIELD"; - case INTERRUPT_ONLY: - return "INTERRUPT_ONLY"; - case ALWAYS_TIME_OUT: - return "ALWAYS_TIME_OUT"; - case ALWAYS_MARK_KILLED: - return "ALWAYS_MARK_KILLED"; - } - MONGO_UNREACHABLE; - } - - static YieldPolicy parseFromBSON(const StringData& element) { - const std::string& yieldPolicy = element.toString(); - if (yieldPolicy == "YIELD_AUTO") { - return YIELD_AUTO; - } - if (yieldPolicy == "WRITE_CONFLICT_RETRY_ONLY") { - return WRITE_CONFLICT_RETRY_ONLY; - } - if (yieldPolicy == "YIELD_MANUAL") { - return YIELD_MANUAL; - } - if (yieldPolicy == "NO_YIELD") { - return NO_YIELD; - } - if (yieldPolicy == "INTERRUPT_ONLY") { - return INTERRUPT_ONLY; - } - if (yieldPolicy == "ALWAYS_TIME_OUT") { - return ALWAYS_TIME_OUT; - } - if (yieldPolicy == "ALWAYS_MARK_KILLED") { - return ALWAYS_MARK_KILLED; - } - MONGO_UNREACHABLE; - } - /** * This class will ensure a PlanExecutor is disposed before it is deleted. */ @@ -247,7 +177,7 @@ public: std::unique_ptr<WorkingSet> ws, std::unique_ptr<PlanStage> rt, const Collection* collection, - YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, NamespaceString nss = NamespaceString(), std::unique_ptr<QuerySolution> qs = nullptr); @@ -263,11 +193,28 @@ public: std::unique_ptr<WorkingSet> ws, std::unique_ptr<PlanStage> rt, const Collection* collection, - YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, NamespaceString nss = NamespaceString(), std::unique_ptr<QuerySolution> qs = nullptr); /** + * These overloads are for SBE. + */ + static StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( + OperationContext* opCtx, + std::unique_ptr<CanonicalQuery> cq, + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root, + NamespaceString nss, + std::unique_ptr<PlanYieldPolicySBE> yieldPolicy); + static StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( + OperationContext* opCtx, + std::unique_ptr<CanonicalQuery> cq, + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root, + NamespaceString nss, + std::queue<std::pair<BSONObj, boost::optional<RecordId>>> stash, + std::unique_ptr<PlanYieldPolicySBE> yieldPolicy); + + /** * A PlanExecutor must be disposed before destruction. In most cases, this will happen * automatically through a PlanExecutor::Deleter or a ClientCursor. */ @@ -482,6 +429,16 @@ public: * for the batch that is currently being built. Otherwise, return an empty object. */ virtual BSONObj getPostBatchResumeToken() const = 0; + + virtual LockPolicy lockPolicy() const = 0; + + /** + * Returns true if this PlanExecutor proxies to a Pipeline of DocumentSources. + * + * TODO SERVER-48478 : Create a new PlanExecutor implementation specifically for executing the + * Pipeline, and delete PipelineProxyStage. + */ + virtual bool isPipelineExecutor() const = 0; }; } // namespace mongo diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp index 596fb4784dc..2f2a6bbc192 100644 --- a/src/mongo/db/query/plan_executor_impl.cpp +++ b/src/mongo/db/query/plan_executor_impl.cpp @@ -53,7 +53,7 @@ #include "mongo/db/exec/working_set_common.h" #include "mongo/db/query/find_common.h" #include "mongo/db/query/mock_yield_policies.h" -#include "mongo/db/query/plan_yield_policy.h" +#include "mongo/db/query/plan_yield_policy_impl.h" #include "mongo/db/repl/replication_coordinator.h" #include "mongo/db/service_context.h" #include "mongo/logv2/log.h" @@ -86,19 +86,19 @@ MONGO_FAIL_POINT_DEFINE(planExecutorHangWhileYieldedInWaitForInserts); * Constructs a PlanYieldPolicy based on 'policy'. */ std::unique_ptr<PlanYieldPolicy> makeYieldPolicy(PlanExecutor* exec, - PlanExecutor::YieldPolicy policy) { + PlanYieldPolicy::YieldPolicy policy) { switch (policy) { - case PlanExecutor::YieldPolicy::YIELD_AUTO: - case PlanExecutor::YieldPolicy::YIELD_MANUAL: - case PlanExecutor::YieldPolicy::NO_YIELD: - case PlanExecutor::YieldPolicy::WRITE_CONFLICT_RETRY_ONLY: - case PlanExecutor::YieldPolicy::INTERRUPT_ONLY: { - return std::make_unique<PlanYieldPolicy>(exec, policy); + case PlanYieldPolicy::YieldPolicy::YIELD_AUTO: + case PlanYieldPolicy::YieldPolicy::YIELD_MANUAL: + case PlanYieldPolicy::YieldPolicy::NO_YIELD: + case PlanYieldPolicy::YieldPolicy::WRITE_CONFLICT_RETRY_ONLY: + case PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY: { + return std::make_unique<PlanYieldPolicyImpl>(exec, policy); } - case PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT: { + case PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT: { return std::make_unique<AlwaysTimeOutYieldPolicy>(exec); } - case PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED: { + case PlanYieldPolicy::YieldPolicy::ALWAYS_MARK_KILLED: { return std::make_unique<AlwaysPlanKilledYieldPolicy>(exec); } default: @@ -132,7 +132,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PlanExecutor::make( std::unique_ptr<WorkingSet> ws, std::unique_ptr<PlanStage> rt, const Collection* collection, - YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, NamespaceString nss, std::unique_ptr<QuerySolution> qs) { auto expCtx = cq->getExpCtx(); @@ -152,7 +152,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PlanExecutor::make( std::unique_ptr<WorkingSet> ws, std::unique_ptr<PlanStage> rt, const Collection* collection, - YieldPolicy yieldPolicy, + PlanYieldPolicy::YieldPolicy yieldPolicy, NamespaceString nss, std::unique_ptr<QuerySolution> qs) { return PlanExecutorImpl::make(expCtx->opCtx, @@ -175,7 +175,7 @@ StatusWith<unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PlanExecutorImpl::ma const boost::intrusive_ptr<ExpressionContext>& expCtx, const Collection* collection, NamespaceString nss, - YieldPolicy yieldPolicy) { + PlanYieldPolicy::YieldPolicy yieldPolicy) { auto execImpl = new PlanExecutorImpl(opCtx, std::move(ws), @@ -206,7 +206,7 @@ PlanExecutorImpl::PlanExecutorImpl(OperationContext* opCtx, const boost::intrusive_ptr<ExpressionContext>& expCtx, const Collection* collection, NamespaceString nss, - YieldPolicy yieldPolicy) + PlanYieldPolicy::YieldPolicy yieldPolicy) : _opCtx(opCtx), _cq(std::move(cq)), _expCtx(_cq ? _cq->getExpCtx() : expCtx), @@ -215,7 +215,8 @@ PlanExecutorImpl::PlanExecutorImpl(OperationContext* opCtx, _root(std::move(rt)), _nss(std::move(nss)), // There's no point in yielding if the collection doesn't exist. - _yieldPolicy(makeYieldPolicy(this, collection ? yieldPolicy : NO_YIELD)) { + _yieldPolicy(makeYieldPolicy( + this, collection ? yieldPolicy : PlanYieldPolicy::YieldPolicy::NO_YIELD)) { invariant(!_expCtx || _expCtx->opCtx == _opCtx); invariant(!_cq || !_expCtx || _cq->getExpCtx() == _expCtx); @@ -336,7 +337,7 @@ void PlanExecutorImpl::restoreState() { throw; // Handles retries by calling restoreStateWithoutRetrying() in a loop. - uassertStatusOK(_yieldPolicy->yieldOrInterrupt()); + uassertStatusOK(_yieldPolicy->yieldOrInterrupt(getOpCtx())); } } @@ -476,7 +477,7 @@ void PlanExecutorImpl::_waitForInserts(CappedInsertNotifierData* notifierData) { ON_BLOCK_EXIT([curOp] { curOp->resumeTimer(); }); auto opCtx = _opCtx; uint64_t currentNotifierVersion = notifierData->notifier->getVersion(); - auto yieldResult = _yieldPolicy->yieldOrInterrupt([opCtx, notifierData] { + auto yieldResult = _yieldPolicy->yieldOrInterrupt(opCtx, [opCtx, notifierData] { const auto deadline = awaitDataState(opCtx).waitForInsertsDeadline; notifierData->notifier->waitUntil(notifierData->lastEOFVersion, deadline); if (MONGO_unlikely(planExecutorHangWhileYieldedInWaitForInserts.shouldFail())) { @@ -527,8 +528,8 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob // 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(_opCtx)) { + uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx)); } WorkingSetID id = WorkingSet::INVALID_ID; @@ -717,4 +718,22 @@ BSONObj PlanExecutorImpl::getPostBatchResumeToken() const { } } +PlanExecutor::LockPolicy PlanExecutorImpl::lockPolicy() const { + if (isPipelineExecutor()) { + return LockPolicy::kLocksInternally; + } + + // If this PlanExecutor is simply unspooling queued data, then there is no need to acquire + // locks. + if (_root->stageType() == StageType::STAGE_QUEUED_DATA) { + return LockPolicy::kLocksInternally; + } + + return LockPolicy::kLockExternally; +} + +bool PlanExecutorImpl::isPipelineExecutor() const { + return _root->stageType() == StageType::STAGE_PIPELINE_PROXY || + _root->stageType() == StageType::STAGE_CHANGE_STREAM_PROXY; +} } // namespace mongo diff --git a/src/mongo/db/query/plan_executor_impl.h b/src/mongo/db/query/plan_executor_impl.h index 0b4e3ca8b24..ac1bcd8e43b 100644 --- a/src/mongo/db/query/plan_executor_impl.h +++ b/src/mongo/db/query/plan_executor_impl.h @@ -53,7 +53,7 @@ public: const boost::intrusive_ptr<ExpressionContext>& expCtx, const Collection* collection, NamespaceString nss, - YieldPolicy yieldPolicy); + PlanYieldPolicy::YieldPolicy yieldPolicy); virtual ~PlanExecutorImpl(); WorkingSet* getWorkingSet() const final; @@ -83,6 +83,8 @@ public: bool isDetached() const final; Timestamp getLatestOplogTimestamp() const final; BSONObj getPostBatchResumeToken() const final; + LockPolicy lockPolicy() const final; + bool isPipelineExecutor() const final; private: /** @@ -96,7 +98,7 @@ private: const boost::intrusive_ptr<ExpressionContext>& expCtx, const Collection* collection, NamespaceString nss, - YieldPolicy yieldPolicy); + PlanYieldPolicy::YieldPolicy yieldPolicy); /** * Clients of PlanExecutor expect that on receiving a new instance from one of the make() diff --git a/src/mongo/db/query/plan_executor_sbe.cpp b/src/mongo/db/query/plan_executor_sbe.cpp new file mode 100644 index 00000000000..187837a1157 --- /dev/null +++ b/src/mongo/db/query/plan_executor_sbe.cpp @@ -0,0 +1,303 @@ +/** + * 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. + */ + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + +#include "mongo/platform/basic.h" + +#include "mongo/db/query/plan_executor_sbe.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/query/sbe_stage_builder.h" +#include "mongo/logv2/log.h" + +namespace mongo { +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PlanExecutor::make( + OperationContext* opCtx, + std::unique_ptr<CanonicalQuery> cq, + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root, + NamespaceString nss, + std::unique_ptr<PlanYieldPolicySBE> yieldPolicy) { + + auto&& [rootStage, data] = root; + + LOGV2_DEBUG(4822860, + 5, + "SBE plan", + "slots"_attr = data.debugString(), + "stages"_attr = sbe::DebugPrinter{}.print(rootStage.get())); + + rootStage->prepare(data.ctx); + + auto exec = new PlanExecutorSBE(opCtx, + std::move(cq), + std::move(root), + std::move(nss), + false, + boost::none, + std::move(yieldPolicy)); + return {{exec, PlanExecutor::Deleter{opCtx}}}; +} + +StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> PlanExecutor::make( + OperationContext* opCtx, + std::unique_ptr<CanonicalQuery> cq, + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root, + NamespaceString nss, + std::queue<std::pair<BSONObj, boost::optional<RecordId>>> stash, + std::unique_ptr<PlanYieldPolicySBE> yieldPolicy) { + + auto&& [rootStage, data] = root; + + LOGV2_DEBUG(4822861, + 5, + "SBE plan", + "slots"_attr = data.debugString(), + "stages"_attr = sbe::DebugPrinter{}.print(rootStage.get())); + + auto exec = new PlanExecutorSBE( + opCtx, std::move(cq), std::move(root), std::move(nss), true, stash, std::move(yieldPolicy)); + return {{exec, PlanExecutor::Deleter{opCtx}}}; +} + +PlanExecutorSBE::PlanExecutorSBE( + OperationContext* opCtx, + std::unique_ptr<CanonicalQuery> cq, + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root, + NamespaceString nss, + bool isOpen, + boost::optional<std::queue<std::pair<BSONObj, boost::optional<RecordId>>>> stash, + std::unique_ptr<PlanYieldPolicySBE> yieldPolicy) + : _state{isOpen ? State::kOpened : State::kClosed}, + _opCtx(opCtx), + _nss(std::move(nss)), + _root(std::move(root.first)), + _cq{std::move(cq)}, + _yieldPolicy(std::move(yieldPolicy)) { + invariant(_root); + + auto&& data = root.second; + + if (data.resultSlot) { + _result = _root->getAccessor(data.ctx, *data.resultSlot); + uassert(4822865, "Query does not have result slot.", _result); + } + + if (data.recordIdSlot) { + _resultRecordId = _root->getAccessor(data.ctx, *data.recordIdSlot); + uassert(4822866, "Query does not have recordId slot.", _resultRecordId); + } + + if (data.oplogTsSlot) { + _oplogTs = _root->getAccessor(data.ctx, *data.oplogTsSlot); + uassert(4822867, "Query does not have oplogTs slot.", _oplogTs); + } + + _shouldTrackLatestOplogTimestamp = data.shouldTrackLatestOplogTimestamp; + _shouldTrackResumeToken = data.shouldTrackResumeToken; + + if (!isOpen) { + _root->attachFromOperationContext(_opCtx); + } + + if (stash) { + _stash = std::move(*stash); + } + + // Callers are allowed to disable yielding for this plan by passing a null yield policy. + if (_yieldPolicy) { + _yieldPolicy->setRootStage(_root.get()); + } +} + +void PlanExecutorSBE::saveState() { + invariant(_root); + _root->saveState(); +} + +void PlanExecutorSBE::restoreState() { + invariant(_root); + _root->restoreState(); +} + +void PlanExecutorSBE::detachFromOperationContext() { + invariant(_opCtx); + invariant(_root); + _root->detachFromOperationContext(); + _opCtx = nullptr; +} + +void PlanExecutorSBE::reattachToOperationContext(OperationContext* opCtx) { + invariant(!_opCtx); + invariant(_root); + _root->attachFromOperationContext(opCtx); + _opCtx = opCtx; +} + +void PlanExecutorSBE::markAsKilled(Status killStatus) { + invariant(!killStatus.isOK()); + // If killed multiple times, only retain the first status. + if (_killStatus.isOK()) { + _killStatus = killStatus; + } +} + +void PlanExecutorSBE::dispose(OperationContext* opCtx) { + if (_root && _state != State::kClosed) { + _root->close(); + _state = State::kClosed; + } + + _root.reset(); +} + +void PlanExecutorSBE::enqueue(const Document& obj) { + enqueue(obj.toBson()); +} + +void PlanExecutorSBE::enqueue(const BSONObj& obj) { + invariant(_state == State::kOpened); + _stash.push({obj.getOwned(), boost::none}); +} + +PlanExecutor::ExecState PlanExecutorSBE::getNext(Document* objOut, RecordId* dlOut) { + invariant(_root); + + BSONObj obj; + auto result = getNext(&obj, dlOut); + if (result == PlanExecutor::ExecState::ADVANCED) { + *objOut = Document{std::move(obj)}; + } + return result; +} + +PlanExecutor::ExecState PlanExecutorSBE::getNext(BSONObj* out, RecordId* dlOut) { + invariant(_root); + + if (!_stash.empty()) { + auto&& [doc, recordId] = _stash.front(); + *out = std::move(doc); + if (dlOut && recordId) { + *dlOut = *recordId; + } + _stash.pop(); + return PlanExecutor::ExecState::ADVANCED; + } else if (_root->getCommonStats()->isEOF) { + // If we had stashed elements and consumed them all, but the PlanStage has also + // already exhausted, we can return EOF straight away. Otherwise, proceed with + // fetching the next document. + _root->close(); + _state = State::kClosed; + return PlanExecutor::ExecState::IS_EOF; + } + + if (_state == State::kClosed) { + _state = State::kOpened; + _root->open(false); + } + + invariant(_state == State::kOpened); + + auto result = fetchNext(_root.get(), _result, _resultRecordId, out, dlOut); + if (result == sbe::PlanState::IS_EOF) { + _root->close(); + _state = State::kClosed; + return PlanExecutor::ExecState::IS_EOF; + } + invariant(result == sbe::PlanState::ADVANCED); + return PlanExecutor::ExecState::ADVANCED; +} + +Timestamp PlanExecutorSBE::getLatestOplogTimestamp() const { + if (_shouldTrackLatestOplogTimestamp) { + invariant(_oplogTs); + + auto [tag, val] = _oplogTs->getViewOfValue(); + uassert(4822868, + "Collection scan was asked to track latest operation time, " + "but found a result without a valid 'ts' field", + tag == sbe::value::TypeTags::Timestamp); + return Timestamp{sbe::value::bitcastTo<uint64_t>(val)}; + } + return {}; +} + +BSONObj PlanExecutorSBE::getPostBatchResumeToken() const { + if (_shouldTrackResumeToken) { + invariant(_resultRecordId); + + auto [tag, val] = _resultRecordId->getViewOfValue(); + uassert(4822869, + "Collection scan was asked to track resume token, " + "but found a result without a valid RecordId", + tag == sbe::value::TypeTags::NumberInt64); + return BSON("$recordId" << sbe::value::bitcastTo<int64_t>(val)); + } + return {}; +} + +sbe::PlanState fetchNext(sbe::PlanStage* root, + sbe::value::SlotAccessor* resultSlot, + sbe::value::SlotAccessor* recordIdSlot, + BSONObj* out, + RecordId* dlOut) { + invariant(out); + + auto state = root->getNext(); + if (state == sbe::PlanState::IS_EOF) { + return state; + } + invariant(state == sbe::PlanState::ADVANCED); + + if (resultSlot) { + auto [tag, val] = resultSlot->getViewOfValue(); + if (tag == sbe::value::TypeTags::Object) { + BSONObjBuilder bb; + sbe::bson::convertToBsonObj(bb, sbe::value::getObjectView(val)); + *out = bb.obj(); + } else if (tag == sbe::value::TypeTags::bsonObject) { + *out = BSONObj(sbe::value::bitcastTo<const char*>(val)); + } else { + // The query is supposed to return an object. + MONGO_UNREACHABLE; + } + } + + if (dlOut) { + invariant(recordIdSlot); + auto [tag, val] = recordIdSlot->getViewOfValue(); + if (tag == sbe::value::TypeTags::NumberInt64) { + *dlOut = RecordId{sbe::value::bitcastTo<int64_t>(val)}; + } + } + return state; +} + +} // namespace mongo diff --git a/src/mongo/db/query/plan_executor_sbe.h b/src/mongo/db/query/plan_executor_sbe.h new file mode 100644 index 00000000000..7d3c169cc84 --- /dev/null +++ b/src/mongo/db/query/plan_executor_sbe.h @@ -0,0 +1,180 @@ +/** + * 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 <queue> + +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/query/plan_executor.h" +#include "mongo/db/query/plan_yield_policy_sbe.h" + +namespace mongo { +class PlanExecutorSBE final : public PlanExecutor { +public: + PlanExecutorSBE( + OperationContext* opCtx, + std::unique_ptr<CanonicalQuery> cq, + std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root, + NamespaceString nss, + bool isOpen, + boost::optional<std::queue<std::pair<BSONObj, boost::optional<RecordId>>>> stash, + std::unique_ptr<PlanYieldPolicySBE> yieldPolicy); + + WorkingSet* getWorkingSet() const override { + MONGO_UNREACHABLE; + } + + PlanStage* getRootStage() const override { + return nullptr; + } + + CanonicalQuery* getCanonicalQuery() const override { + return _cq.get(); + } + + const NamespaceString& nss() const override { + return _nss; + } + + OperationContext* getOpCtx() const override { + return _opCtx; + } + + const boost::intrusive_ptr<ExpressionContext>& getExpCtx() const override { + static boost::intrusive_ptr<ExpressionContext> unused; + return unused; + } + + void saveState(); + void restoreState(); + + void detachFromOperationContext(); + void reattachToOperationContext(OperationContext* opCtx); + + void restoreStateWithoutRetrying() override { + MONGO_UNREACHABLE; + } + + ExecState getNextSnapshotted(Snapshotted<Document>* objOut, RecordId* dlOut) override { + MONGO_UNREACHABLE; + } + + ExecState getNextSnapshotted(Snapshotted<BSONObj>* objOut, RecordId* dlOut) override { + MONGO_UNREACHABLE; + } + + ExecState getNext(Document* objOut, RecordId* dlOut) override; + ExecState getNext(BSONObj* out, RecordId* dlOut) override; + + bool isEOF() override { + return _state == State::kClosed; + } + + void executePlan() override { + MONGO_UNREACHABLE; + } + + void markAsKilled(Status killStatus); + + void dispose(OperationContext* opCtx); + + void enqueue(const Document& obj); + void enqueue(const BSONObj& obj); + + bool isMarkedAsKilled() const override { + return !_killStatus.isOK(); + } + + Status getKillStatus() override { + invariant(isMarkedAsKilled()); + return _killStatus; + } + + bool isDisposed() const override { + return !_root; + } + + bool isDetached() const override { + return !_opCtx; + } + + Timestamp getLatestOplogTimestamp() const override; + BSONObj getPostBatchResumeToken() const override; + + LockPolicy lockPolicy() const override { + return LockPolicy::kLocksInternally; + } + + bool isPipelineExecutor() const override { + return false; + } + +private: + enum class State { kClosed, kOpened }; + + State _state{State::kClosed}; + + OperationContext* _opCtx; + + NamespaceString _nss; + + std::unique_ptr<sbe::PlanStage> _root; + + sbe::value::SlotAccessor* _result{nullptr}; + sbe::value::SlotAccessor* _resultRecordId{nullptr}; + sbe::value::SlotAccessor* _oplogTs{nullptr}; + bool _shouldTrackLatestOplogTimestamp{false}; + bool _shouldTrackResumeToken{false}; + + std::queue<std::pair<BSONObj, boost::optional<RecordId>>> _stash; + + // If _killStatus has a non-OK value, then we have been killed and the value represents the + // reason for the kill. + Status _killStatus = Status::OK(); + + std::unique_ptr<CanonicalQuery> _cq; + + std::unique_ptr<PlanYieldPolicySBE> _yieldPolicy; +}; + +/** + * Executes getNext() on the 'root' PlanStage and used 'resultSlot' and 'recordIdSlot' to access the + * fetched document and it's record id, which are stored in 'out' and 'dlOut' parameters + * respectively, if they not null pointers. + * + * This common logic can be used by various consumers which need to fetch data using an SBE + * PlanStage tree, such as PlanExecutor or RuntimePlanner. + */ +sbe::PlanState fetchNext(sbe::PlanStage* root, + sbe::value::SlotAccessor* resultSlot, + sbe::value::SlotAccessor* recordIdSlot, + BSONObj* out, + RecordId* dlOut); +} // namespace mongo diff --git a/src/mongo/db/query/plan_ranker.cpp b/src/mongo/db/query/plan_ranker.cpp index e81612cf7e7..fa47db3f5e5 100644 --- a/src/mongo/db/query/plan_ranker.cpp +++ b/src/mongo/db/query/plan_ranker.cpp @@ -31,279 +31,129 @@ #include "mongo/platform/basic.h" -#include <algorithm> -#include <cmath> -#include <utility> -#include <vector> - #include "mongo/db/query/plan_ranker.h" -#include "mongo/db/exec/plan_stage.h" -#include "mongo/db/exec/working_set.h" -#include "mongo/db/query/explain.h" -#include "mongo/db/query/query_knobs_gen.h" -#include "mongo/db/query/query_solution.h" -#include "mongo/db/server_options.h" #include "mongo/logv2/log.h" -namespace { - -/** - * Comparator for (scores, candidateIndex) in pickBestPlan(). - */ -bool scoreComparator(const std::pair<double, size_t>& lhs, const std::pair<double, size_t>& rhs) { - // Just compare score in lhs.first and rhs.first; - // Ignore candidate array index in lhs.second and rhs.second. - return lhs.first > rhs.first; +namespace mongo::plan_ranker { +namespace log_detail { +void logScoreFormula(std::function<std::string()> formula, + double score, + double baseScore, + double productivity, + double noFetchBonus, + double noSortBonus, + double noIxisectBonus, + double tieBreakers) { + LOGV2_DEBUG(20961, 2, "{sb_str}", "sb_str"_attr = [&]() { + StringBuilder sb; + sb << "score(" << str::convertDoubleToString(score) << ") = baseScore(" + << str::convertDoubleToString(baseScore) << ")" + << " + productivity(" << formula() << " = " << str::convertDoubleToString(productivity) + << ")" + << " + tieBreakers(" << str::convertDoubleToString(noFetchBonus) << " noFetchBonus + " + << str::convertDoubleToString(noSortBonus) << " noSortBonus + " + << str::convertDoubleToString(noIxisectBonus) + << " noIxisectBonus = " << str::convertDoubleToString(tieBreakers) << ")"; + return sb.str(); + }()); } -} // namespace - -namespace mongo { - -using std::endl; -using std::vector; - -// static -StatusWith<std::unique_ptr<PlanRankingDecision>> PlanRanker::pickBestPlan( - const vector<CandidatePlan>& candidates) { - invariant(!candidates.empty()); - // A plan that hits EOF is automatically scored above - // its peers. If multiple plans hit EOF during the same - // set of round-robin calls to work(), then all such plans - // receive the bonus. - double eofBonus = 1.0; - - // Each plan will have a stat tree. - std::vector<std::unique_ptr<PlanStageStats>> statTrees; - - // Get stat trees from each plan. - // Copy stats trees instead of transferring ownership - // because multi plan runner will need its own stats - // trees for explain. - for (size_t i = 0; i < candidates.size(); ++i) { - statTrees.push_back(candidates[i].root->getStats()); - } - - // Holds (score, candidateInndex). - // Used to derive scores and candidate ordering. - vector<std::pair<double, size_t>> scoresAndCandidateindices; - vector<size_t> failed; +void logScoreBoost(double score) { + LOGV2_DEBUG(20962, + 5, + "Score boosted to {newScore} due to intersection forcing", + "Score boosted due to intersection forcing", + "newScore"_attr = score); +} - // Compute score for each tree. Record the best. - for (size_t i = 0; i < statTrees.size(); ++i) { - if (!candidates[i].failed) { - LOGV2_DEBUG( - 20956, +void logScoringPlan(std::function<std::string()> solution, + std::function<std::string()> explain, + std::function<std::string()> planSummary, + size_t planIndex, + bool isEOF) { + LOGV2_DEBUG(20956, 5, "Scoring plan {planIndex}:\n{querySolution}Stats:\n{stats}", "Scoring plan", - "planIndex"_attr = i, - "querySolution"_attr = redact(candidates[i].solution->toString()), - "stats"_attr = redact( - Explain::statsToBSON(*statTrees[i]).jsonString(ExtendedRelaxedV2_0_0, true))); - LOGV2_DEBUG(20957, - 2, - "Scoring query plan: {planSummary} planHitEOF={planHitEOF}", - "Scoring query plan", - "planSummary"_attr = Explain::getPlanSummary(candidates[i].root), - "planHitEOF"_attr = statTrees[i]->common.isEOF); - - double score = scoreTree(statTrees[i].get()); - LOGV2_DEBUG( - 20958, 5, "Basic plan score: {score}", "Basic plan score", "score"_attr = score); - if (statTrees[i]->common.isEOF) { - LOGV2_DEBUG(20959, - 5, - "Adding +{eofBonus} EOF bonus to score", - "Adding EOF bonus to score", - "eofBonus"_attr = eofBonus); - score += 1; - } - - scoresAndCandidateindices.push_back(std::make_pair(score, i)); - } else { - failed.push_back(i); - LOGV2_DEBUG(20960, - 2, - "Not scoring plan: {planSummary} because the plan failed", - "Not scoring a plan because the plan failed", - "planSummary"_attr = Explain::getPlanSummary(candidates[i].root)); - } - } - - // If there isn't a viable plan we should error. - if (scoresAndCandidateindices.size() == 0U) { - return {ErrorCodes::Error(31157), - "No viable plan was found because all candidate plans failed."}; - } - - // Sort (scores, candidateIndex). Get best child and populate candidate ordering. - std::stable_sort( - scoresAndCandidateindices.begin(), scoresAndCandidateindices.end(), scoreComparator); - - auto why = std::make_unique<PlanRankingDecision>(); - - // Determine whether plans tied for the win. - if (scoresAndCandidateindices.size() > 1U) { - double bestScore = scoresAndCandidateindices[0].first; - double runnerUpScore = scoresAndCandidateindices[1].first; - const double epsilon = 1e-10; - why->tieForBest = std::abs(bestScore - runnerUpScore) < epsilon; - } - - // Update results in 'why' - // Stats and scores in 'why' are sorted in descending order by score. - why->stats.clear(); - why->scores.clear(); - why->candidateOrder.clear(); - why->failedCandidates = std::move(failed); - for (size_t i = 0; i < scoresAndCandidateindices.size(); ++i) { - double score = scoresAndCandidateindices[i].first; - size_t candidateIndex = scoresAndCandidateindices[i].second; - - // We shouldn't cache the scores with the EOF bonus included, - // as this is just a tie-breaking measure for plan selection. - // Plans not run through the multi plan runner will not receive - // the bonus. - // - // An example of a bad thing that could happen if we stored scores - // with the EOF bonus included: - // - // Let's say Plan A hits EOF, is the highest ranking plan, and gets - // cached as such. On subsequent runs it will not receive the bonus. - // Eventually the plan cache feedback mechanism will evict the cache - // entry---the scores will appear to have fallen due to the missing - // EOF bonus. - // - // This begs the question, why don't we include the EOF bonus in - // scoring of cached plans as well? The problem here is that the cached - // plan runner always runs plans to completion before scoring. Queries - // that don't get the bonus in the multi plan runner might get the bonus - // after being run from the plan cache. - if (statTrees[candidateIndex]->common.isEOF) { - score -= eofBonus; - } - - why->stats.push_back(std::move(statTrees[candidateIndex])); - why->scores.push_back(score); - why->candidateOrder.push_back(candidateIndex); - } - for (auto& i : why->failedCandidates) { - why->stats.push_back(std::move(statTrees[i])); - } - - return StatusWith<std::unique_ptr<PlanRankingDecision>>(std::move(why)); + "planIndex"_attr = planIndex, + "querySolution"_attr = [solution]() { return redact(solution()); }(), + "stats"_attr = [explain]() { return redact(explain()); }()); + LOGV2_DEBUG(20957, + 2, + "Scoring query plan: {planSummary} planHitEOF={planHitEOF}", + "Scoring query plan", + "planSummary"_attr = [planSummary]() { return planSummary(); }(), + "planHitEOF"_attr = isEOF); } -// TODO: Move this out. This is a signal for ranking but will become its own complicated -// stats-collecting beast. -double computeSelectivity(const PlanStageStats* stats) { - if (STAGE_IXSCAN == stats->stageType) { - IndexScanStats* iss = static_cast<IndexScanStats*>(stats->specific.get()); - return iss->keyPattern.nFields(); - } else { - double sum = 0; - for (size_t i = 0; i < stats->children.size(); ++i) { - sum += computeSelectivity(stats->children[i].get()); - } - return sum; - } +void logScore(double score) { + LOGV2_DEBUG(20958, 5, "Basic plan score: {score}", "Basic plan score", "score"_attr = score); } -bool hasStage(const StageType type, const PlanStageStats* stats) { - if (type == stats->stageType) { - return true; - } - for (size_t i = 0; i < stats->children.size(); ++i) { - if (hasStage(type, stats->children[i].get())) { - return true; - } - } - return false; +void logEOFBonus(double eofBonus) { + LOGV2_DEBUG(20959, + 5, + "Adding +{eofBonus} EOF bonus to score", + "Adding EOF bonus to score", + "eofBonus"_attr = eofBonus); } -// static -double PlanRanker::scoreTree(const PlanStageStats* stats) { - // We start all scores at 1. Our "no plan selected" score is 0 and we want all plans to - // be greater than that. - double baseScore = 1; - - // How many "units of work" did the plan perform. Each call to work(...) - // counts as one unit. - size_t workUnits = stats->common.works; - invariant(workUnits != 0); - - // How much did a plan produce? - // Range: [0, 1] - double productivity = - static_cast<double>(stats->common.advanced) / static_cast<double>(workUnits); - - // Just enough to break a tie. Must be small enough to ensure that a more productive - // plan doesn't lose to a less productive plan due to tie breaking. - const double epsilon = std::min(1.0 / static_cast<double>(10 * workUnits), 1e-4); +void logFailedPlan(std::function<std::string()> planSummary) { + LOGV2_DEBUG(20960, + 2, + "Not scoring plan: {planSummary} because the plan failed", + "Not scoring a plan because the plan failed", + "planSummary"_attr = [&]() { return planSummary(); }()); +} +} // namespace log_detail - // We prefer queries that don't require a fetch stage. - double noFetchBonus = epsilon; - if (hasStage(STAGE_FETCH, stats)) { - noFetchBonus = 0; +namespace { +/** + * A plan scorer for the classic plan stage tree. Defines the plan productivity as a number + * of intermediate results returned, or advanced, by the root stage, divided by the "unit of works" + * which the plan performed. Each call to work(...) counts as one unit. + */ +class DefaultPlanScorer final : public PlanScorer<PlanStageStats> { +protected: + double calculateProductivity(const PlanStageStats* stats) const final { + invariant(stats->common.works != 0); + return static_cast<double>(stats->common.advanced) / + static_cast<double>(stats->common.works); } - // In the case of ties, prefer solutions without a blocking sort - // to solutions with a blocking sort. - double noSortBonus = epsilon; - if (hasStage(STAGE_SORT_DEFAULT, stats) || hasStage(STAGE_SORT_SIMPLE, stats)) { - noSortBonus = 0; + std::string getProductivityFormula(const PlanStageStats* stats) const { + StringBuilder sb; + sb << "(" << stats->common.advanced << " advanced)/(" << stats->common.works << " works)"; + return sb.str(); } - // In the case of ties, prefer single index solutions to ixisect. Index - // intersection solutions are often slower than single-index solutions - // because they require examining a superset of index keys that would be - // examined by a single index scan. - // - // On the other hand, index intersection solutions examine the same - // number or fewer of documents. In the case that index intersection - // allows us to examine fewer documents, the penalty given to ixisect - // can be made up via the no fetch bonus. - double noIxisectBonus = epsilon; - if (hasStage(STAGE_AND_HASH, stats) || hasStage(STAGE_AND_SORTED, stats)) { - noIxisectBonus = 0; + double getNumberOfAdvances(const PlanStageStats* stats) const final { + return stats->common.works; } - double tieBreakers = noFetchBonus + noSortBonus + noIxisectBonus; - double score = baseScore + productivity + tieBreakers; + bool hasStage(StageType type, const PlanStageStats* root) const final { + std::queue<const PlanStageStats*> remaining; + remaining.push(root); - if (shouldLog(logv2::LogSeverity::Debug(2))) { - StringBuilder sb; - sb << "baseScore(" << str::convertDoubleToString(baseScore) << ")" - << " + productivity((" << stats->common.advanced << " advanced)/(" << stats->common.works - << " works) = " << str::convertDoubleToString(productivity) << ")" - << " + tieBreakers(" << str::convertDoubleToString(noFetchBonus) << " noFetchBonus + " - << str::convertDoubleToString(noSortBonus) << " noSortBonus + " - << str::convertDoubleToString(noIxisectBonus) - << " noIxisectBonus = " << str::convertDoubleToString(tieBreakers) << ")"; - LOGV2_DEBUG(20961, - 2, - "score({score}) = {calculation}", - "Plan score calculation", - "score"_attr = score, - "calculation"_attr = sb.str()); - } + while (!remaining.empty()) { + auto stats = remaining.front(); + remaining.pop(); + + if (stats->stageType == type) { + return true; + } - if (internalQueryForceIntersectionPlans.load()) { - if (hasStage(STAGE_AND_HASH, stats) || hasStage(STAGE_AND_SORTED, stats)) { - // The boost should be >2.001 to make absolutely sure the ixisect plan will win due - // to the combination of 1) productivity, 2) eof bonus, and 3) no ixisect bonus. - score += 3; - LOGV2_DEBUG(20962, - 5, - "Score boosted to {newScore} due to intersection forcing", - "Score boosted due to intersection forcing", - "newScore"_attr = score); + for (auto&& child : stats->children) { + remaining.push(child.get()); + } } + return false; } +}; +} // namespace - return score; +std::unique_ptr<PlanScorer<PlanStageStats>> makePlanScorer() { + return std::make_unique<DefaultPlanScorer>(); } - -} // namespace mongo +} // namespace mongo::plan_ranker diff --git a/src/mongo/db/query/plan_ranker.h b/src/mongo/db/query/plan_ranker.h index 581a99a59c9..d9df7c8557a 100644 --- a/src/mongo/db/query/plan_ranker.h +++ b/src/mongo/db/query/plan_ranker.h @@ -29,62 +29,170 @@ #pragma once -#include <memory> #include <queue> -#include <vector> -#include "mongo/base/owned_pointer_vector.h" -#include "mongo/db/exec/plan_stage.h" #include "mongo/db/exec/plan_stats.h" +#include "mongo/db/exec/sbe/stages/plan_stats.h" #include "mongo/db/exec/working_set.h" +#include "mongo/db/query/explain.h" #include "mongo/db/query/query_solution.h" #include "mongo/util/container_size_helper.h" -namespace mongo { - -struct CandidatePlan; -struct PlanRankingDecision; +namespace mongo::plan_ranker { +// The logging facility enforces the rule that logging should not be done in a header file. Since +// template classes and functions below must be defined in the header file and since they use the +// logging facility, we have to define the helper functions below to perform the actual logging +// operation from template code. +// Note that we pass std::function callback instead of string values to avoid spending time +// generating log output that may never actually be written to the logs, depending on the current +// log level. +namespace log_detail { +void logScoreFormula(std::function<std::string()> formula, + double score, + double baseScore, + double productivity, + double noFetchBonus, + double noSortBonus, + double noIxisectBonus, + double tieBreakers); +void logScoreBoost(double score); +void logScoringPlan(std::function<std::string()> solution, + std::function<std::string()> explain, + std::function<std::string()> planSummary, + size_t planIndex, + bool isEOF); +void logScore(double score); +void logEOFBonus(double eofBonus); +void logFailedPlan(std::function<std::string()> planSummary); +} // namespace log_detail /** - * Ranks 2 or more plans. + * Assigns the stats tree a 'goodness' score. The higher the score, the better the plan. The exact + * value isn't meaningful except for imposing a ranking. + * + * All specific plan scorers should inherit from this scorer and provide methods to produce the plan + * productivity factor, and the number of plan "advances", representing the number of documents + * returned by the PlanStage tree. */ -class PlanRanker { +template <typename PlanStageStatsType> +class PlanScorer { public: + PlanScorer() = default; + virtual ~PlanScorer() = default; + + double calculateScore(const PlanStageStatsType* stats) const { + // We start all scores at 1. Our "no plan selected" score is 0 and we want all plans to + // be greater than that. + const double baseScore = 1; + + const auto productivity = calculateProductivity(stats); + const auto advances = getNumberOfAdvances(stats); + const double epsilon = + std::min(1.0 / static_cast<double>(10 * (advances > 0 ? advances : 1)), 1e-4); + + + // We prefer queries that don't require a fetch stage. + double noFetchBonus = epsilon; + if (hasStage(STAGE_FETCH, stats)) { + noFetchBonus = 0; + } + + // In the case of ties, prefer solutions without a blocking sort + // to solutions with a blocking sort. + double noSortBonus = epsilon; + if (hasStage(STAGE_SORT_DEFAULT, stats) || hasStage(STAGE_SORT_SIMPLE, stats)) { + noSortBonus = 0; + } + + // In the case of ties, prefer single index solutions to ixisect. Index + // intersection solutions are often slower than single-index solutions + // because they require examining a superset of index keys that would be + // examined by a single index scan. + // + // On the other hand, index intersection solutions examine the same + // number or fewer of documents. In the case that index intersection + // allows us to examine fewer documents, the penalty given to ixisect + // can be made up via the no fetch bonus. + double noIxisectBonus = epsilon; + if (hasStage(STAGE_AND_HASH, stats) || hasStage(STAGE_AND_SORTED, stats)) { + noIxisectBonus = 0; + } + + const double tieBreakers = noFetchBonus + noSortBonus + noIxisectBonus; + double score = baseScore + productivity + tieBreakers; + + log_detail::logScoreFormula([this, stats] { return getProductivityFormula(stats); }, + score, + baseScore, + productivity, + noFetchBonus, + noSortBonus, + noIxisectBonus, + tieBreakers); + + if (internalQueryForceIntersectionPlans.load()) { + if (hasStage(STAGE_AND_HASH, stats) || hasStage(STAGE_AND_SORTED, stats)) { + // The boost should be >2.001 to make absolutely sure the ixisect plan will win due + // to the combination of 1) productivity, 2) eof bonus, and 3) no ixisect bonus. + score += 3; + log_detail::logScoreBoost(score); + } + } + return score; + } + +protected: + /** + * Returns an abstract plan productivity value. Each implementation is free to define the + * formula to calculate the productivity. The value must be withing the range: [0, 1]. + */ + virtual double calculateProductivity(const PlanStageStatsType* stats) const = 0; + + /** + * Returns a string desribing a formula to calculte plan producivity. It can be used for the log + * output, for example. + */ + virtual std::string getProductivityFormula(const PlanStageStatsType* stats) const = 0; + /** - * Returns a PlanRankingDecision which has the ranking and the information about the ranking - * process with status OK if everything worked. 'candidateOrder' within the PlanRankingDecision - * holds indices into candidates ordered by score (winner in first element). - * - * Returns an error if there was an issue with plan ranking (e.g. there was no viable plan). + * Returns the number of advances from the root stage stats, which represents the number of + * documents returned by the PlanStage tree. */ - static StatusWith<std::unique_ptr<PlanRankingDecision>> pickBestPlan( - const std::vector<CandidatePlan>& candidates); + virtual double getNumberOfAdvances(const PlanStageStatsType* stats) const = 0; /** - * Assign the stats tree a 'goodness' score. The higher the score, the better - * the plan. The exact value isn't meaningful except for imposing a ranking. + * True, if the plan stage stats tree represents a plan stage of the given 'type'. */ - static double scoreTree(const PlanStageStats* stats); + virtual bool hasStage(StageType type, const PlanStageStatsType* stats) const = 0; }; /** - * A container holding one to-be-ranked plan and its associated/relevant data. - * Does not own any of its pointers. + * A container holding one to-be-scored plan and its associated/relevant data. + * It takes the following template parameters: + * * PlanStageType - the type of plan stages in the execution tree. + * * ResultType - the type of data produced by the execution tree during the candidate plan + * execution. + * * Data - the type of any auxiliary data which is needed to run the execution tree. */ -struct CandidatePlan { - CandidatePlan(std::unique_ptr<QuerySolution> solution, PlanStage* r, WorkingSet* w) - : solution(std::move(solution)), root(r), ws(w), failed(false) {} - +template <typename PlanStageType, typename ResultType, typename Data> +struct BaseCandidatePlan { + // A query solution representing this candidate plan. std::unique_ptr<QuerySolution> solution; - PlanStage* root; // Not owned here. - WorkingSet* ws; // Not owned here. - - // Any results produced during the plan's execution prior to ranking are retained here. - std::queue<WorkingSetID> results; - - bool failed; + // A root stage of the PlanStage tree constructed from the 'solution'. + PlanStageType root; + // Any auxiliary data required to run the execution tree. + Data data; + // Indicates whether this candidate plan has completed the trial run early by achieving one + // of the trial run metrics. + bool exitedEarly{false}; + // Indicates that this candidate plan has failed in a recoverable fashion during the trial run. + bool failed{false}; + // Any results produced during the plan's execution prior to scoring are retained here. + std::queue<ResultType> results; }; +using CandidatePlan = BaseCandidatePlan<PlanStage*, WorkingSetID, WorkingSet*>; + /** * Information about why a plan was picked to be the best. Data here is placed into the cache * and used to compare expected performance with actual. @@ -97,11 +205,17 @@ struct PlanRankingDecision { */ PlanRankingDecision* clone() const { PlanRankingDecision* decision = new PlanRankingDecision(); - for (size_t i = 0; i < stats.size(); ++i) { - PlanStageStats* s = stats[i].get(); - invariant(s); - decision->stats.push_back(std::unique_ptr<PlanStageStats>{s->clone()}); - } + stdx::visit( + [&decision](auto&& planStats) { + using StatsType = typename std::decay_t<decltype(planStats)>::value_type; + std::vector<StatsType> copy; + for (auto&& stats : planStats) { + invariant(stats); + copy.push_back(StatsType{stats->clone()}); + } + decision->stats = std::move(copy); + }, + stats); decision->scores = scores; decision->candidateOrder = candidateOrder; decision->failedCandidates = failedCandidates; @@ -110,8 +224,12 @@ struct PlanRankingDecision { uint64_t estimateObjectSizeInBytes() const { return // Add size of each element in 'stats' vector. - container_size_helper::estimateObjectSizeInBytes( - stats, [](const auto& stat) { return stat->estimateObjectSizeInBytes(); }, true) + + stdx::visit( + [](auto&& stats) { + return container_size_helper::estimateObjectSizeInBytes( + stats, [](auto&& stat) { return stat->estimateObjectSizeInBytes(); }, true); + }, + stats) + // Add size of each element in 'candidateOrder' vector. container_size_helper::estimateObjectSizeInBytes(candidateOrder) + // Add size of each element in 'failedCandidates' vector. @@ -122,16 +240,27 @@ struct PlanRankingDecision { sizeof(*this); } + template <typename PlanStageStatsType> + const std::vector<std::unique_ptr<PlanStageStatsType>>& getStats() const { + return stdx::get<std::vector<std::unique_ptr<PlanStageStatsType>>>(stats); + } + + template <typename PlanStageStatsType> + std::vector<std::unique_ptr<PlanStageStatsType>>& getStats() { + return stdx::get<std::vector<std::unique_ptr<PlanStageStatsType>>>(stats); + } + // Stats of all plans sorted in descending order by score. - // Owned by us. - std::vector<std::unique_ptr<PlanStageStats>> stats; + stdx::variant<std::vector<std::unique_ptr<PlanStageStats>>, + std::vector<std::unique_ptr<mongo::sbe::PlanStageStats>>> + stats; // The "goodness" score corresponding to 'stats'. // Sorted in descending order. std::vector<double> scores; // Ordering of original plans in descending of score. - // Filled in by PlanRanker::pickBestPlan(candidates, ...) + // Filled in by PlanScorer::pickBestPlan(candidates, ...) // so that candidates[candidateOrder[0]] refers to the best plan // with corresponding cores[0] and stats[0]. Runner-up would be // candidates[candidateOrder[1]] followed by @@ -152,4 +281,143 @@ struct PlanRankingDecision { bool tieForBest = false; }; -} // namespace mongo +/** + * A factory function to create a plan scorer for a plan stage stats tree. + */ +std::unique_ptr<PlanScorer<PlanStageStats>> makePlanScorer(); +} // namespace mongo::plan_ranker + +// Forward declaration. +namespace mongo::sbe::plan_ranker { +std::unique_ptr<mongo::plan_ranker::PlanScorer<PlanStageStats>> makePlanScorer( + const QuerySolution* solution); +} // namespace mongo::sbe::plan_ranker + +namespace mongo::plan_ranker { +/** + * Returns a PlanRankingDecision which has the ranking and the information about the ranking + * process with status OK if everything worked. 'candidateOrder' within the PlanRankingDecision + * holds indices into candidates ordered by score (winner in first element). + * + * Returns an error if there was an issue with plan ranking (e.g. there was no viable plan). + */ +template <typename PlanStageStatsType, typename PlanStageType, typename ResultType, typename Data> +StatusWith<std::unique_ptr<PlanRankingDecision>> pickBestPlan( + const std::vector<BaseCandidatePlan<PlanStageType, ResultType, Data>>& candidates) { + invariant(!candidates.empty()); + // A plan that hits EOF is automatically scored above + // its peers. If multiple plans hit EOF during the same + // set of round-robin calls to work(), then all such plans + // receive the bonus. + double eofBonus = 1.0; + + // Get stat trees from each plan. + std::vector<std::unique_ptr<PlanStageStatsType>> statTrees; + for (size_t i = 0; i < candidates.size(); ++i) { + statTrees.push_back(candidates[i].root->getStats()); + } + + // Holds (score, candidateIndex). + // Used to derive scores and candidate ordering. + std::vector<std::pair<double, size_t>> scoresAndCandidateIndices; + std::vector<size_t> failed; + + // Compute score for each tree. Record the best. + for (size_t i = 0; i < statTrees.size(); ++i) { + if (!candidates[i].failed) { + log_detail::logScoringPlan( + [& candidate = candidates[i]]() { return candidate.solution->toString(); }, + [& stats = *statTrees[i]]() { + return Explain::statsToBSON(stats).jsonString(ExtendedRelaxedV2_0_0, true); + }, + [root = &*candidates[i].root]() { return Explain::getPlanSummary(root); }, + i, + statTrees[i]->common.isEOF); + auto scorer = [solution = candidates[i].solution.get()]() + -> std::unique_ptr<PlanScorer<PlanStageStatsType>> { + if constexpr (std::is_same_v<PlanStageStatsType, PlanStageStats>) { + return makePlanScorer(); + } else { + static_assert(std::is_same_v<PlanStageStatsType, mongo::sbe::PlanStageStats>); + return sbe::plan_ranker::makePlanScorer(solution); + } + }(); + double score = scorer->calculateScore(statTrees[i].get()); + log_detail::logScore(score); + if (statTrees[i]->common.isEOF) { + log_detail::logEOFBonus(eofBonus); + score += 1; + } + + scoresAndCandidateIndices.push_back(std::make_pair(score, i)); + } else { + failed.push_back(i); + log_detail::logFailedPlan( + [root = &*candidates[i].root] { return Explain::getPlanSummary(root); }); + } + } + + // If there isn't a viable plan we should error. + if (scoresAndCandidateIndices.size() == 0U) { + return {ErrorCodes::Error(31157), + "No viable plan was found because all candidate plans failed."}; + } + + // Sort (scores, candidateIndex). Get best child and populate candidate ordering. + std::stable_sort(scoresAndCandidateIndices.begin(), + scoresAndCandidateIndices.end(), + [](const auto& lhs, const auto& rhs) { + // Just compare score in lhs.first and rhs.first; + // Ignore candidate array index in lhs.second and rhs.second. + return lhs.first > rhs.first; + }); + + auto why = std::make_unique<PlanRankingDecision>(); + why->stats = std::vector<std::unique_ptr<PlanStageStatsType>>{}; + + // Determine whether plans tied for the win. + if (scoresAndCandidateIndices.size() > 1U) { + double bestScore = scoresAndCandidateIndices[0].first; + double runnerUpScore = scoresAndCandidateIndices[1].first; + const double epsilon = 1e-10; + why->tieForBest = std::abs(bestScore - runnerUpScore) < epsilon; + } + + // Update results in 'why' + // Stats and scores in 'why' are sorted in descending order by score. + why->failedCandidates = std::move(failed); + for (size_t i = 0; i < scoresAndCandidateIndices.size(); ++i) { + double score = scoresAndCandidateIndices[i].first; + size_t candidateIndex = scoresAndCandidateIndices[i].second; + + // We shouldn't cache the scores with the EOF bonus included, as this is just a + // tie-breaking measure for plan selection. Plans not run through the multi plan runner + // will not receive the bonus. + // + // An example of a bad thing that could happen if we stored scores with the EOF bonus + // included: + // + // Let's say Plan A hits EOF, is the highest ranking plan, and gets cached as such. On + // subsequent runs it will not receive the bonus. Eventually the plan cache feedback + // mechanism will evict the cache entry - the scores will appear to have fallen due to + // the missing EOF bonus. + // + // This raises the question, why don't we include the EOF bonus in scoring of cached plans + // as well? The problem here is that the cached plan runner always runs plans to completion + // before scoring. Queries that don't get the bonus in the multi plan runner might get the + // bonus after being run from the plan cache. + if (statTrees[candidateIndex]->common.isEOF) { + score -= eofBonus; + } + + why->getStats<PlanStageStatsType>().push_back(std::move(statTrees[candidateIndex])); + why->scores.push_back(score); + why->candidateOrder.push_back(candidateIndex); + } + for (auto& i : why->failedCandidates) { + why->getStats<PlanStageStatsType>().push_back(std::move(statTrees[i])); + } + + return StatusWith<std::unique_ptr<PlanRankingDecision>>(std::move(why)); +} +} // namespace mongo::plan_ranker diff --git a/src/mongo/db/query/plan_ranker_test.cpp b/src/mongo/db/query/plan_ranker_test.cpp index 8f51a0edab1..229b7e4fae3 100644 --- a/src/mongo/db/query/plan_ranker_test.cpp +++ b/src/mongo/db/query/plan_ranker_test.cpp @@ -70,8 +70,9 @@ TEST(PlanRankerTest, NoFetchBonus) { badPlan->children[0]->children.emplace_back( makeStats("IXSCAN", STAGE_IXSCAN, make_unique<IndexScanStats>())); - auto goodScore = PlanRanker::scoreTree(goodPlan.get()); - auto badScore = PlanRanker::scoreTree(badPlan.get()); + auto scorer = plan_ranker::makePlanScorer(); + auto goodScore = scorer->calculateScore(goodPlan.get()); + auto badScore = scorer->calculateScore(badPlan.get()); ASSERT_GT(goodScore, badScore); } diff --git a/src/mongo/db/query/plan_yield_policy.cpp b/src/mongo/db/query/plan_yield_policy.cpp index 448ac827797..56010f83708 100644 --- a/src/mongo/db/query/plan_yield_policy.cpp +++ b/src/mongo/db/query/plan_yield_policy.cpp @@ -31,48 +31,26 @@ #include "mongo/db/query/plan_yield_policy.h" -#include "mongo/db/concurrency/write_conflict_exception.h" -#include "mongo/db/curop.h" -#include "mongo/db/curop_failpoint_helpers.h" #include "mongo/db/operation_context.h" -#include "mongo/db/query/query_knobs_gen.h" -#include "mongo/db/service_context.h" -#include "mongo/db/storage/snapshot_helper.h" #include "mongo/util/fail_point.h" #include "mongo/util/scopeguard.h" #include "mongo/util/time_support.h" namespace mongo { -namespace { -MONGO_FAIL_POINT_DEFINE(setInterruptOnlyPlansCheckForInterruptHang); -} // namespace +PlanYieldPolicy::PlanYieldPolicy(YieldPolicy policy, + ClockSource* cs, + int yieldIterations, + Milliseconds yieldPeriod) + : _policy(policy), _elapsedTracker(cs, yieldIterations, yieldPeriod) {} -PlanYieldPolicy::PlanYieldPolicy(PlanExecutor* exec, PlanExecutor::YieldPolicy policy) - : _policy(exec->getOpCtx()->lockState()->isGlobalLockedRecursively() ? PlanExecutor::NO_YIELD - : policy), - _forceYield(false), - _elapsedTracker(exec->getOpCtx()->getServiceContext()->getFastClockSource(), - internalQueryExecYieldIterations.load(), - Milliseconds(internalQueryExecYieldPeriodMS.load())), - _planYielding(exec) {} - - -PlanYieldPolicy::PlanYieldPolicy(PlanExecutor::YieldPolicy policy, ClockSource* cs) - : _policy(policy), - _forceYield(false), - _elapsedTracker(cs, - internalQueryExecYieldIterations.load(), - Milliseconds(internalQueryExecYieldPeriodMS.load())), - _planYielding(nullptr) {} - -bool PlanYieldPolicy::shouldYieldOrInterrupt() { - if (_policy == PlanExecutor::INTERRUPT_ONLY) { +bool PlanYieldPolicy::shouldYieldOrInterrupt(OperationContext* opCtx) { + if (_policy == YieldPolicy::INTERRUPT_ONLY) { return _elapsedTracker.intervalHasElapsed(); } if (!canAutoYield()) return false; - invariant(!_planYielding->getOpCtx()->lockState()->inAWriteUnitOfWork()); + invariant(!opCtx->lockState()->inAWriteUnitOfWork()); if (_forceYield) return true; return _elapsedTracker.intervalHasElapsed(); @@ -82,143 +60,27 @@ void PlanYieldPolicy::resetTimer() { _elapsedTracker.resetLastTime(); } -Status PlanYieldPolicy::yieldOrInterrupt(std::function<void()> whileYieldingFn) { - invariant(_planYielding); +Status PlanYieldPolicy::yieldOrInterrupt(OperationContext* opCtx, + std::function<void()> whileYieldingFn) { + invariant(opCtx); - if (_policy == PlanExecutor::INTERRUPT_ONLY) { + if (_policy == YieldPolicy::INTERRUPT_ONLY) { ON_BLOCK_EXIT([this]() { resetTimer(); }); - OperationContext* opCtx = _planYielding->getOpCtx(); invariant(opCtx); - // If the 'setInterruptOnlyPlansCheckForInterruptHang' fail point is enabled, set the 'msg' - // field of this operation's CurOp to signal that we've hit this point. - if (MONGO_unlikely(setInterruptOnlyPlansCheckForInterruptHang.shouldFail())) { - CurOpFailpointHelpers::waitWhileFailPointEnabled( - &setInterruptOnlyPlansCheckForInterruptHang, - opCtx, - "setInterruptOnlyPlansCheckForInterruptHang"); - } - + preCheckInterruptOnly(opCtx); return opCtx->checkForInterruptNoAssert(); } - invariant(canAutoYield()); + + invariant(!opCtx->lockState()->inAWriteUnitOfWork()); // After we finish yielding (or in any early return), call resetTimer() to prevent yielding // again right away. We delay the resetTimer() call so that the clock doesn't start ticking // until after we return from the yield. ON_BLOCK_EXIT([this]() { resetTimer(); }); - _forceYield = false; - OperationContext* opCtx = _planYielding->getOpCtx(); - invariant(opCtx); - invariant(!opCtx->lockState()->inAWriteUnitOfWork()); - - // Can't use writeConflictRetry since we need to call saveState before reseting the transaction. - for (int attempt = 1; true; attempt++) { - try { - try { - _planYielding->saveState(); - } catch (const WriteConflictException&) { - invariant(!"WriteConflictException not allowed in saveState"); - } - - if (_policy == PlanExecutor::WRITE_CONFLICT_RETRY_ONLY) { - // Just reset the snapshot. Leave all LockManager locks alone. - opCtx->recoveryUnit()->abandonSnapshot(); - } else { - // Release and reacquire locks. - _yieldAllLocks(opCtx, whileYieldingFn, _planYielding->nss()); - } - - _planYielding->restoreStateWithoutRetrying(); - return Status::OK(); - } catch (const WriteConflictException&) { - CurOp::get(opCtx)->debug().additiveMetrics.incrementWriteConflicts(1); - WriteConflictException::logAndBackoff( - attempt, "plan execution restoreState", _planYielding->nss().ns()); - // retry - } catch (...) { - // Errors other than write conflicts don't get retried, and should instead result in the - // PlanExecutor dying. We propagate all such errors as status codes. - return exceptionToStatus(); - } - } -} - -namespace { -MONGO_FAIL_POINT_DEFINE(setYieldAllLocksHang); -MONGO_FAIL_POINT_DEFINE(setYieldAllLocksWait); -} // namespace - -void PlanYieldPolicy::_yieldAllLocks(OperationContext* opCtx, - std::function<void()> whileYieldingFn, - const NamespaceString& planExecNS) { - // Things have to happen here in a specific order: - // * Release lock mgr locks - // * Check for interrupt (kill flag is set) - // * Call the whileYieldingFn - // * Reacquire lock mgr locks - - Locker* locker = opCtx->lockState(); - - Locker::LockSnapshot snapshot; - - auto unlocked = locker->saveLockStateAndUnlock(&snapshot); - - // Attempt to check for interrupt while locks are not held, in order to discourage the - // assumption that locks will always be held when a Plan Executor returns an error. - if (_policy == PlanExecutor::YIELD_AUTO) { - opCtx->checkForInterrupt(); // throws - } - - if (!unlocked) { - // Nothing was unlocked, just return, yielding is pointless. - return; - } - - // Top-level locks are freed, release any potential low-level (storage engine-specific - // locks). If we are yielding, we are at a safe place to do so. - opCtx->recoveryUnit()->abandonSnapshot(); - - // Track the number of yields in CurOp. - CurOp::get(opCtx)->yielded(); - - setYieldAllLocksHang.executeIf( - [opCtx](const BSONObj& config) { - setYieldAllLocksHang.pauseWhileSet(); - - if (config.getField("checkForInterruptAfterHang").trueValue()) { - // Throws. - opCtx->checkForInterrupt(); - } - }, - [&](const BSONObj& config) { - StringData ns = config.getStringField("namespace"); - return ns.empty() || ns == planExecNS.ns(); - }); - setYieldAllLocksWait.executeIf( - [&](const BSONObj& data) { sleepFor(Milliseconds(data["waitForMillis"].numberInt())); }, - [&](const BSONObj& config) { - BSONElement dataNs = config["namespace"]; - return !dataNs || planExecNS.ns() == dataNs.str(); - }); - - if (whileYieldingFn) { - whileYieldingFn(); - } - - locker->restoreLockState(opCtx, snapshot); - - // After yielding and reacquiring locks, the preconditions that were used to select our - // ReadSource initially need to be checked again. Queries hold an AutoGetCollectionForRead RAII - // lock for their lifetime, which may select a ReadSource based on state (e.g. replication - // state). After a query yields its locks, this state may have changed, invalidating our current - // choice of ReadSource. Using the same preconditions, change our ReadSource if necessary. - auto newReadSource = SnapshotHelper::getNewReadSource(opCtx, planExecNS); - if (newReadSource) { - opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource); - } + return yield(opCtx, whileYieldingFn); } } // namespace mongo diff --git a/src/mongo/db/query/plan_yield_policy.h b/src/mongo/db/query/plan_yield_policy.h index 512bc908dcc..3c8a85cbbd7 100644 --- a/src/mongo/db/query/plan_yield_policy.h +++ b/src/mongo/db/query/plan_yield_policy.h @@ -31,8 +31,8 @@ #include <functional> -#include "mongo/db/catalog/collection.h" -#include "mongo/db/query/plan_executor.h" +#include "mongo/db/operation_context.h" +#include "mongo/util/duration.h" #include "mongo/util/elapsed_tracker.h" namespace mongo { @@ -41,23 +41,116 @@ class ClockSource; class PlanYieldPolicy { public: - virtual ~PlanYieldPolicy() {} + enum class YieldPolicy { + // Any call to getNext() may yield. In particular, the executor may die on any call to + // getNext() due to a required index or collection becoming invalid during yield. If this + // occurs, getNext() will produce an error during yield recovery and will return FAILURE. + // Additionally, this will handle all WriteConflictExceptions that occur while processing + // the query. With this yield policy, it is possible for getNext() to return FAILURE with + // locks released, if the operation is killed while yielding. + YIELD_AUTO, + + // This will handle WriteConflictExceptions that occur while processing the query, but will + // not yield locks. abandonSnapshot() will be called if a WriteConflictException occurs so + // callers must be prepared to get a new snapshot. The caller must hold their locks + // continuously from construction to destruction. Callers which do not want auto-yielding, + // but may release their locks during query execution must use the YIELD_MANUAL policy. + WRITE_CONFLICT_RETRY_ONLY, + + // Use this policy if you want to disable auto-yielding, but will release locks while using + // the PlanExecutor. Any WriteConflictExceptions will be raised to the caller of getNext(). + // + // With this policy, an explicit call must be made to saveState() before releasing locks, + // and an explicit call to restoreState() must be made after reacquiring locks. + // restoreState() will throw if the PlanExecutor is now invalid due to a catalog operation + // (e.g. collection drop) during yield. + YIELD_MANUAL, + + // Can be used in one of the following scenarios: + // - The caller will hold a lock continuously for the lifetime of this PlanExecutor. + // - This PlanExecutor doesn't logically belong to a Collection, and so does not need to be + // locked during execution. For example, a PlanExecutor containing a PipelineProxyStage + // which is being used to execute an aggregation pipeline. + NO_YIELD, + + // Will not yield locks or storage engine resources, but will check for interrupt. + INTERRUPT_ONLY, + + // Used for testing, this yield policy will cause the PlanExecutor to time out on the first + // yield, returning FAILURE with an error object encoding a ErrorCodes::ExceededTimeLimit + // message. + ALWAYS_TIME_OUT, + + // Used for testing, this yield policy will cause the PlanExecutor to be marked as killed on + // the first yield, returning FAILURE with an error object encoding a + // ErrorCodes::QueryPlanKilled message. + ALWAYS_MARK_KILLED, + }; + + static std::string serializeYieldPolicy(YieldPolicy yieldPolicy) { + switch (yieldPolicy) { + case YieldPolicy::YIELD_AUTO: + return "YIELD_AUTO"; + case YieldPolicy::WRITE_CONFLICT_RETRY_ONLY: + return "WRITE_CONFLICT_RETRY_ONLY"; + case YieldPolicy::YIELD_MANUAL: + return "YIELD_MANUAL"; + case YieldPolicy::NO_YIELD: + return "NO_YIELD"; + case YieldPolicy::INTERRUPT_ONLY: + return "INTERRUPT_ONLY"; + case YieldPolicy::ALWAYS_TIME_OUT: + return "ALWAYS_TIME_OUT"; + case YieldPolicy::ALWAYS_MARK_KILLED: + return "ALWAYS_MARK_KILLED"; + } + MONGO_UNREACHABLE; + } - PlanYieldPolicy(PlanExecutor* exec, PlanExecutor::YieldPolicy policy); + static YieldPolicy parseFromBSON(const StringData& element) { + const std::string& yieldPolicy = element.toString(); + if (yieldPolicy == "YIELD_AUTO") { + return YieldPolicy::YIELD_AUTO; + } + if (yieldPolicy == "WRITE_CONFLICT_RETRY_ONLY") { + return YieldPolicy::WRITE_CONFLICT_RETRY_ONLY; + } + if (yieldPolicy == "YIELD_MANUAL") { + return YieldPolicy::YIELD_MANUAL; + } + if (yieldPolicy == "NO_YIELD") { + return YieldPolicy::NO_YIELD; + } + if (yieldPolicy == "INTERRUPT_ONLY") { + return YieldPolicy::INTERRUPT_ONLY; + } + if (yieldPolicy == "ALWAYS_TIME_OUT") { + return YieldPolicy::ALWAYS_TIME_OUT; + } + if (yieldPolicy == "ALWAYS_MARK_KILLED") { + return YieldPolicy::ALWAYS_MARK_KILLED; + } + MONGO_UNREACHABLE; + } /** - * Only used in dbtests since we don't have access to a PlanExecutor. Since we don't have - * access to the PlanExecutor to grab a ClockSource from, we pass in a ClockSource directly - * in the constructor instead. + * Constructs a PlanYieldPolicy of the given 'policy' type. This class uses an ElapsedTracker + * to keep track of elapsed time, which is initialized from the parameters 'cs', + * 'yieldIterations' and 'yieldPeriod'. */ - PlanYieldPolicy(PlanExecutor::YieldPolicy policy, ClockSource* cs); + PlanYieldPolicy(YieldPolicy policy, + ClockSource* cs, + int yieldIterations, + Milliseconds yieldPeriod); + + virtual ~PlanYieldPolicy() = default; /** * Periodically returns true to indicate that it is time to check for interrupt (in the case of * YIELD_AUTO and INTERRUPT_ONLY) or release locks or storage engine state (in the case of * auto-yielding plans). */ - virtual bool shouldYieldOrInterrupt(); + virtual bool shouldYieldOrInterrupt(OperationContext* opCtx); /** * Resets the yield timer so that we wait for a while before yielding/interrupting again. @@ -77,7 +170,8 @@ public: * Calls 'whileYieldingFn' after relinquishing locks and before reacquiring the locks that have * been relinquished. */ - virtual Status yieldOrInterrupt(std::function<void()> whileYieldingFn = nullptr); + Status yieldOrInterrupt(OperationContext* opCtx, + std::function<void()> whileYieldingFn = nullptr); /** * All calls to shouldYieldOrInterrupt() will return true until the next call to @@ -95,15 +189,15 @@ public: */ bool canReleaseLocksDuringExecution() const { switch (_policy) { - case PlanExecutor::YIELD_AUTO: - case PlanExecutor::YIELD_MANUAL: - case PlanExecutor::ALWAYS_TIME_OUT: - case PlanExecutor::ALWAYS_MARK_KILLED: { + case YieldPolicy::YIELD_AUTO: + case YieldPolicy::YIELD_MANUAL: + case YieldPolicy::ALWAYS_TIME_OUT: + case YieldPolicy::ALWAYS_MARK_KILLED: { return true; } - case PlanExecutor::NO_YIELD: - case PlanExecutor::WRITE_CONFLICT_RETRY_ONLY: - case PlanExecutor::INTERRUPT_ONLY: { + case YieldPolicy::NO_YIELD: + case YieldPolicy::WRITE_CONFLICT_RETRY_ONLY: + case YieldPolicy::INTERRUPT_ONLY: { return false; } } @@ -117,45 +211,41 @@ public: */ bool canAutoYield() const { switch (_policy) { - case PlanExecutor::YIELD_AUTO: - case PlanExecutor::WRITE_CONFLICT_RETRY_ONLY: - case PlanExecutor::ALWAYS_TIME_OUT: - case PlanExecutor::ALWAYS_MARK_KILLED: { + case YieldPolicy::YIELD_AUTO: + case YieldPolicy::WRITE_CONFLICT_RETRY_ONLY: + case YieldPolicy::ALWAYS_TIME_OUT: + case YieldPolicy::ALWAYS_MARK_KILLED: { return true; } - case PlanExecutor::NO_YIELD: - case PlanExecutor::YIELD_MANUAL: - case PlanExecutor::INTERRUPT_ONLY: + case YieldPolicy::NO_YIELD: + case YieldPolicy::YIELD_MANUAL: + case YieldPolicy::INTERRUPT_ONLY: return false; } MONGO_UNREACHABLE; } - PlanExecutor::YieldPolicy getPolicy() const { + PlanYieldPolicy::YieldPolicy getPolicy() const { return _policy; } private: - const PlanExecutor::YieldPolicy _policy; - - bool _forceYield; - ElapsedTracker _elapsedTracker; - - // The plan executor which this yield policy is responsible for yielding. Must - // not outlive the plan executor. - PlanExecutor* const _planYielding; + /** + * Yields locks and calls 'abandonSnapshot()'. Calls 'whileYieldingFn()', if provided, while + * locks are not held. + */ + virtual Status yield(OperationContext* opCtx, + std::function<void()> whileYieldingFn = nullptr) = 0; /** - * If not in a nested context, unlocks all locks, suggests to the operating system to - * switch to another thread, and then reacquires all locks. - * - * If in a nested context (eg DBDirectClient), does nothing. - * - * The whileYieldingFn will be executed after unlocking the locks and before re-acquiring them. + * If the yield policy is INTERRUPT_ONLY, this is called prior to checking for interrupt. */ - void _yieldAllLocks(OperationContext* opCtx, - std::function<void()> whileYieldingFn, - const NamespaceString& planExecNS); + virtual void preCheckInterruptOnly(OperationContext* opCtx) {} + + const YieldPolicy _policy; + + bool _forceYield = false; + ElapsedTracker _elapsedTracker; }; } // namespace mongo diff --git a/src/mongo/db/query/plan_yield_policy_impl.cpp b/src/mongo/db/query/plan_yield_policy_impl.cpp new file mode 100644 index 00000000000..3f6a0b0f9ff --- /dev/null +++ b/src/mongo/db/query/plan_yield_policy_impl.cpp @@ -0,0 +1,176 @@ +/** + * 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/query/plan_yield_policy_impl.h" + +#include "mongo/db/concurrency/write_conflict_exception.h" +#include "mongo/db/curop.h" +#include "mongo/db/curop_failpoint_helpers.h" +#include "mongo/db/query/query_knobs_gen.h" +#include "mongo/db/service_context.h" +#include "mongo/db/storage/snapshot_helper.h" +#include "mongo/util/fail_point.h" + +namespace mongo { +namespace { +MONGO_FAIL_POINT_DEFINE(setInterruptOnlyPlansCheckForInterruptHang); +} // namespace + +PlanYieldPolicyImpl::PlanYieldPolicyImpl(PlanExecutor* exec, PlanYieldPolicy::YieldPolicy policy) + : PlanYieldPolicy(exec->getOpCtx()->lockState()->isGlobalLockedRecursively() + ? PlanYieldPolicy::YieldPolicy::NO_YIELD + : policy, + exec->getOpCtx()->getServiceContext()->getFastClockSource(), + internalQueryExecYieldIterations.load(), + Milliseconds{internalQueryExecYieldPeriodMS.load()}), + _planYielding(exec) {} + +Status PlanYieldPolicyImpl::yield(OperationContext* opCtx, std::function<void()> whileYieldingFn) { + // Can't use writeConflictRetry since we need to call saveState before reseting the + // transaction. + for (int attempt = 1; true; attempt++) { + try { + try { + _planYielding->saveState(); + } catch (const WriteConflictException&) { + invariant(!"WriteConflictException not allowed in saveState"); + } + + if (getPolicy() == PlanYieldPolicy::YieldPolicy::WRITE_CONFLICT_RETRY_ONLY) { + // Just reset the snapshot. Leave all LockManager locks alone. + opCtx->recoveryUnit()->abandonSnapshot(); + } else { + // Release and reacquire locks. + _yieldAllLocks(opCtx, whileYieldingFn, _planYielding->nss()); + } + + _planYielding->restoreStateWithoutRetrying(); + return Status::OK(); + } catch (const WriteConflictException&) { + CurOp::get(opCtx)->debug().additiveMetrics.incrementWriteConflicts(1); + WriteConflictException::logAndBackoff( + attempt, "plan execution restoreState", _planYielding->nss().ns()); + // retry + } catch (...) { + // Errors other than write conflicts don't get retried, and should instead result in + // the PlanExecutor dying. We propagate all such errors as status codes. + return exceptionToStatus(); + } + } +} + +namespace { +MONGO_FAIL_POINT_DEFINE(setYieldAllLocksHang); +MONGO_FAIL_POINT_DEFINE(setYieldAllLocksWait); +} // namespace + +void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx, + std::function<void()> whileYieldingFn, + const NamespaceString& planExecNS) { + // Things have to happen here in a specific order: + // * Release lock mgr locks + // * Check for interrupt (kill flag is set) + // * Call the whileYieldingFn + // * Reacquire lock mgr locks + + Locker* locker = opCtx->lockState(); + + Locker::LockSnapshot snapshot; + + auto unlocked = locker->saveLockStateAndUnlock(&snapshot); + + // Attempt to check for interrupt while locks are not held, in order to discourage the + // assumption that locks will always be held when a Plan Executor returns an error. + if (getPolicy() == PlanYieldPolicy::YieldPolicy::YIELD_AUTO) { + opCtx->checkForInterrupt(); // throws + } + + if (!unlocked) { + // Nothing was unlocked, just return, yielding is pointless. + return; + } + + // Top-level locks are freed, release any potential low-level (storage engine-specific + // locks). If we are yielding, we are at a safe place to do so. + opCtx->recoveryUnit()->abandonSnapshot(); + + // Track the number of yields in CurOp. + CurOp::get(opCtx)->yielded(); + + setYieldAllLocksHang.executeIf( + [opCtx](const BSONObj& config) { + setYieldAllLocksHang.pauseWhileSet(); + + if (config.getField("checkForInterruptAfterHang").trueValue()) { + // Throws. + opCtx->checkForInterrupt(); + } + }, + [&](const BSONObj& config) { + StringData ns = config.getStringField("namespace"); + return ns.empty() || ns == planExecNS.ns(); + }); + setYieldAllLocksWait.executeIf( + [&](const BSONObj& data) { sleepFor(Milliseconds(data["waitForMillis"].numberInt())); }, + [&](const BSONObj& config) { + BSONElement dataNs = config["namespace"]; + return !dataNs || planExecNS.ns() == dataNs.str(); + }); + + if (whileYieldingFn) { + whileYieldingFn(); + } + + locker->restoreLockState(opCtx, snapshot); + + // After yielding and reacquiring locks, the preconditions that were used to select our + // ReadSource initially need to be checked again. Queries hold an AutoGetCollectionForRead RAII + // lock for their lifetime, which may select a ReadSource based on state (e.g. replication + // state). After a query yields its locks, this state may have changed, invalidating our current + // choice of ReadSource. Using the same preconditions, change our ReadSource if necessary. + auto newReadSource = SnapshotHelper::getNewReadSource(opCtx, planExecNS); + if (newReadSource) { + opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource); + } +} + +void PlanYieldPolicyImpl::preCheckInterruptOnly(OperationContext* opCtx) { + // If the 'setInterruptOnlyPlansCheckForInterruptHang' fail point is enabled, set the + // 'failPointMsg' field of this operation's CurOp to signal that we've hit this point. + if (MONGO_unlikely(setInterruptOnlyPlansCheckForInterruptHang.shouldFail())) { + CurOpFailpointHelpers::waitWhileFailPointEnabled( + &setInterruptOnlyPlansCheckForInterruptHang, + opCtx, + "setInterruptOnlyPlansCheckForInterruptHang"); + } +} + +} // namespace mongo diff --git a/src/mongo/db/query/plan_yield_policy_impl.h b/src/mongo/db/query/plan_yield_policy_impl.h new file mode 100644 index 00000000000..fcc92669fe7 --- /dev/null +++ b/src/mongo/db/query/plan_yield_policy_impl.h @@ -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. + */ + +#pragma once + +#include "mongo/db/query/plan_executor.h" +#include "mongo/db/query/plan_yield_policy.h" + +namespace mongo { + +class PlanYieldPolicyImpl final : public PlanYieldPolicy { +public: + PlanYieldPolicyImpl(PlanExecutor* exec, PlanYieldPolicy::YieldPolicy policy); + +private: + Status yield(OperationContext* opCtx, std::function<void()> whileYieldingFn = nullptr) override; + + void preCheckInterruptOnly(OperationContext* opCtx) override; + + /** + * If not in a nested context, unlocks all locks, suggests to the operating system to switch to + * another thread, and then reacquires all locks. + * + * If in a nested context (eg DBDirectClient), does nothing. + * + * The whileYieldingFn will be executed after unlocking the locks and before re-acquiring them. + */ + void _yieldAllLocks(OperationContext* opCtx, + std::function<void()> whileYieldingFn, + const NamespaceString& planExecNS); + + // The plan executor which this yield policy is responsible for yielding. Must not outlive the + // plan executor. + PlanExecutor* const _planYielding; +}; + +} // namespace mongo diff --git a/src/mongo/db/query/plan_yield_policy_sbe.cpp b/src/mongo/db/query/plan_yield_policy_sbe.cpp new file mode 100644 index 00000000000..c7a76154f3c --- /dev/null +++ b/src/mongo/db/query/plan_yield_policy_sbe.cpp @@ -0,0 +1,61 @@ +/** + * 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/query/plan_yield_policy_sbe.h" + +namespace mongo { + +// TODO SERVER-48616: Need to change how we bump the CurOp yield counter. We shouldn't do it here +// so that SBE does not depend on CurOp. But we should still expose that statistic and keep it fresh +// as the operation executes. +Status PlanYieldPolicySBE::yield(OperationContext* opCtx, std::function<void()> whileYieldingFn) { + if (!_rootStage) { + // This yield policy isn't bound to an execution tree yet. + return Status::OK(); + } + + try { + _rootStage->saveState(); + + opCtx->recoveryUnit()->abandonSnapshot(); + + if (whileYieldingFn) { + whileYieldingFn(); + } + + _rootStage->restoreState(); + } catch (...) { + return exceptionToStatus(); + } + + return Status::OK(); +} +} // namespace mongo diff --git a/src/mongo/db/query/plan_yield_policy_sbe.h b/src/mongo/db/query/plan_yield_policy_sbe.h new file mode 100644 index 00000000000..1b9041bcc0a --- /dev/null +++ b/src/mongo/db/query/plan_yield_policy_sbe.h @@ -0,0 +1,59 @@ +/** + * 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/plan_yield_policy.h" + +namespace mongo { + +class PlanYieldPolicySBE final : public PlanYieldPolicy { +public: + PlanYieldPolicySBE(YieldPolicy policy, + ClockSource* clockSource, + int yieldFrequency, + Milliseconds yieldPeriod) + : PlanYieldPolicy(policy, clockSource, yieldFrequency, yieldPeriod) { + uassert(4822879, + "WRITE_CONFLICT_RETRY_ONLY yield policy is not supported in SBE", + policy != YieldPolicy::WRITE_CONFLICT_RETRY_ONLY); + } + + void setRootStage(sbe::PlanStage* rootStage) { + _rootStage = rootStage; + } + +private: + Status yield(OperationContext* opCtx, std::function<void()> whileYieldingFn = nullptr) override; + + sbe::PlanStage* _rootStage = nullptr; +}; + +} // namespace mongo diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp index c9e41f6a7e2..105dba7e8d3 100644 --- a/src/mongo/db/query/planner_analysis.cpp +++ b/src/mongo/db/query/planner_analysis.cpp @@ -253,20 +253,6 @@ void replaceNodeInTree(QuerySolutionNode** root, } } -bool hasNode(QuerySolutionNode* root, StageType type) { - if (type == root->getType()) { - return true; - } - - for (size_t i = 0; i < root->children.size(); ++i) { - if (hasNode(root->children[i], type)) { - return true; - } - } - - return false; -} - void geoSkipValidationOn(const std::set<StringData>& twoDSphereFields, QuerySolutionNode* solnRoot) { // If there is a GeoMatchExpression in the tree on a field with a 2dsphere index, @@ -892,7 +878,7 @@ std::unique_ptr<QuerySolution> QueryPlannerAnalysis::analyzeDataAccess( // A solution can be blocking if it has a blocking sort stage or // a hashed AND stage. - bool hasAndHashStage = hasNode(solnRoot.get(), STAGE_AND_HASH); + bool hasAndHashStage = solnRoot->hasNode(STAGE_AND_HASH); soln->hasBlockingStage = hasSortStage || hasAndHashStage; const QueryRequest& qr = query.getQueryRequest(); diff --git a/src/mongo/db/query/projection.cpp b/src/mongo/db/query/projection.cpp index 0943c5538d8..5c9c26c0abe 100644 --- a/src/mongo/db/query/projection.cpp +++ b/src/mongo/db/query/projection.cpp @@ -31,7 +31,7 @@ #include "mongo/base/exact_cast.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 { @@ -174,7 +174,7 @@ auto analyzeProjection(const ProjectionPathASTNode* root, ProjectType type) { ProjectionAnalysisVisitor projectionAnalysisVisitor{&deps}; PathTrackingWalker walker{&context, {&depsAnalysisVisitor, &projectionAnalysisVisitor}, {}}; - projection_ast_walker::walk(&walker, root); + tree_walker::walk<true, projection_ast::ASTNode>(root, &walker); const auto& userData = context.data(); const auto& tracker = userData.fieldDependencyTracker; diff --git a/src/mongo/db/query/projection_ast.h b/src/mongo/db/query/projection_ast.h index 3f85c85c419..3f68b7ff34c 100644 --- a/src/mongo/db/query/projection_ast.h +++ b/src/mongo/db/query/projection_ast.h @@ -98,6 +98,22 @@ protected: ASTNodeVector _children; }; +inline auto begin(ASTNode& node) { + return node.children().begin(); +} + +inline auto begin(const ASTNode& node) { + return node.children().begin(); +} + +inline auto end(ASTNode& node) { + return node.children().end(); +} + +inline auto end(const ASTNode& node) { + return node.children().end(); +} + class MatchExpressionASTNode final : public ASTNode { public: MatchExpressionASTNode(CopyableMatchExpression matchExpr) : _matchExpr{matchExpr} {} diff --git a/src/mongo/db/query/projection_ast_path_tracking_visitor.h b/src/mongo/db/query/projection_ast_path_tracking_visitor.h index 3909b8ad892..8b33ff1af93 100644 --- a/src/mongo/db/query/projection_ast_path_tracking_visitor.h +++ b/src/mongo/db/query/projection_ast_path_tracking_visitor.h @@ -117,7 +117,7 @@ public: invariant(_context); } - void visit(MaybeConstPtr<IsConst, ProjectionPathASTNode> node) final { + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionPathASTNode> node) final { if (node->parent()) { _context->setBasePath(_context->fullPath()); _context->popFrontFieldName(); @@ -126,12 +126,12 @@ public: _context->pushFieldNames({node->fieldNames().begin(), node->fieldNames().end()}); } - void visit(MaybeConstPtr<IsConst, MatchExpressionASTNode> node) final {} - void visit(MaybeConstPtr<IsConst, ProjectionPositionalASTNode> node) final {} - void visit(MaybeConstPtr<IsConst, ProjectionSliceASTNode> node) final {} - void visit(MaybeConstPtr<IsConst, ProjectionElemMatchASTNode> node) final {} - void visit(MaybeConstPtr<IsConst, ExpressionASTNode> node) final {} - void visit(MaybeConstPtr<IsConst, BooleanConstantASTNode> node) final {} + void visit(tree_walker::MaybeConstPtr<IsConst, MatchExpressionASTNode> node) final {} + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionPositionalASTNode> node) final {} + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionSliceASTNode> node) final {} + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionElemMatchASTNode> node) final {} + void visit(tree_walker::MaybeConstPtr<IsConst, ExpressionASTNode> node) final {} + void visit(tree_walker::MaybeConstPtr<IsConst, BooleanConstantASTNode> node) final {} private: PathTrackingVisitorContext<UserData>* _context; @@ -150,7 +150,7 @@ public: invariant(_context); } - void visit(MaybeConstPtr<IsConst, ProjectionPathASTNode> node) final { + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionPathASTNode> node) final { _context->popFieldNames(); if (_context->basePath()) { @@ -165,27 +165,27 @@ public: } } - void visit(MaybeConstPtr<IsConst, ProjectionPositionalASTNode> node) final { + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionPositionalASTNode> node) final { _context->popFrontFieldName(); } - void visit(MaybeConstPtr<IsConst, ProjectionSliceASTNode> node) final { + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionSliceASTNode> node) final { _context->popFrontFieldName(); } - void visit(MaybeConstPtr<IsConst, ProjectionElemMatchASTNode> node) final { + void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionElemMatchASTNode> node) final { _context->popFrontFieldName(); } - void visit(MaybeConstPtr<IsConst, ExpressionASTNode> node) final { + void visit(tree_walker::MaybeConstPtr<IsConst, ExpressionASTNode> node) final { _context->popFrontFieldName(); } - void visit(MaybeConstPtr<IsConst, BooleanConstantASTNode> node) final { + void visit(tree_walker::MaybeConstPtr<IsConst, BooleanConstantASTNode> node) final { _context->popFrontFieldName(); } - void visit(MaybeConstPtr<IsConst, MatchExpressionASTNode> node) final {} + void visit(tree_walker::MaybeConstPtr<IsConst, MatchExpressionASTNode> node) final {} private: PathTrackingVisitorContext<UserData>* _context; @@ -225,19 +225,19 @@ public: _postVisitors.push_back(&_pathTrackingPostVisitor); } - void preVisit(MaybeConstPtr<IsConst, projection_ast::ASTNode> node) { + void preVisit(tree_walker::MaybeConstPtr<IsConst, projection_ast::ASTNode> node) { for (auto visitor : _preVisitors) { node->acceptVisitor(visitor); } } - void postVisit(MaybeConstPtr<IsConst, projection_ast::ASTNode> node) { + void postVisit(tree_walker::MaybeConstPtr<IsConst, projection_ast::ASTNode> node) { for (auto visitor : _postVisitors) { node->acceptVisitor(visitor); } } - void inVisit(long count, MaybeConstPtr<IsConst, ASTNode> node) {} + void inVisit(long count, tree_walker::MaybeConstPtr<IsConst, ASTNode> node) {} private: PathTrackingPreVisitor<UserData, IsConst> _pathTrackingPreVisitor; diff --git a/src/mongo/db/query/projection_ast_util.cpp b/src/mongo/db/query/projection_ast_util.cpp index ec2ec49560a..e5b4cc1a9c4 100644 --- a/src/mongo/db/query/projection_ast_util.cpp +++ b/src/mongo/db/query/projection_ast_util.cpp @@ -32,7 +32,7 @@ #include "mongo/db/query/projection_ast_util.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" namespace mongo::projection_ast { namespace { @@ -128,7 +128,7 @@ BSONObj astToDebugBSON(const ASTNode* root) { BSONPostVisitor postVisitor{&context.data()}; PathTrackingWalker walker{&context, {&preVisitor}, {&postVisitor}}; - projection_ast_walker::walk(&walker, root); + tree_walker::walk<true, projection_ast::ASTNode>(root, &walker); invariant(context.data().builders.size() == 1); return context.data().builders.top().obj(); diff --git a/src/mongo/db/query/projection_ast_visitor.h b/src/mongo/db/query/projection_ast_visitor.h index 81cc07bf932..30dd8b9df43 100644 --- a/src/mongo/db/query/projection_ast_visitor.h +++ b/src/mongo/db/query/projection_ast_visitor.h @@ -29,6 +29,8 @@ #pragma once +#include "mongo/db/query/tree_walker.h" + namespace mongo { namespace projection_ast { class MatchExpressionASTNode; @@ -40,13 +42,6 @@ class ExpressionASTNode; class BooleanConstantASTNode; /** - * A template type which resolves to 'const T*' if 'IsConst' argument is 'true', and to 'T*' - * otherwise. - */ -template <bool IsConst, typename T> -using MaybeConstPtr = typename std::conditional<IsConst, const T*, T*>::type; - -/** * Visitor pattern for ProjectionAST. * * This code is not responsible for traversing the AST, only for performing the double-dispatch. @@ -58,13 +53,13 @@ template <bool IsConst = false> class ProjectionASTVisitor { public: virtual ~ProjectionASTVisitor() = default; - virtual void visit(MaybeConstPtr<IsConst, MatchExpressionASTNode> node) = 0; - virtual void visit(MaybeConstPtr<IsConst, ProjectionPathASTNode> node) = 0; - virtual void visit(MaybeConstPtr<IsConst, ProjectionPositionalASTNode> node) = 0; - virtual void visit(MaybeConstPtr<IsConst, ProjectionSliceASTNode> node) = 0; - virtual void visit(MaybeConstPtr<IsConst, ProjectionElemMatchASTNode> node) = 0; - virtual void visit(MaybeConstPtr<IsConst, ExpressionASTNode> node) = 0; - virtual void visit(MaybeConstPtr<IsConst, BooleanConstantASTNode> node) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, MatchExpressionASTNode> node) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionPathASTNode> node) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionPositionalASTNode> node) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionSliceASTNode> node) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ProjectionElemMatchASTNode> node) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, ExpressionASTNode> node) = 0; + virtual void visit(tree_walker::MaybeConstPtr<IsConst, BooleanConstantASTNode> node) = 0; }; using ProjectionASTMutableVisitor = ProjectionASTVisitor<false>; diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl index 618e9aedd8d..3035d6ed545 100644 --- a/src/mongo/db/query/query_knobs.idl +++ b/src/mongo/db/query/query_knobs.idl @@ -346,3 +346,21 @@ server_parameters: cpp_varname: "internalQueryDesugarWhereToFunction" cpp_vartype: AtomicWord<bool> default: false + + internalQueryEnableSlotBasedExecutionEngine: + description: "If true, activates the slot-based execution engine to execute queries." + set_at: [ startup, runtime ] + cpp_varname: "internalQueryEnableSlotBasedExecutionEngine" + cpp_vartype: AtomicWord<bool> + default: false + + internalQueryDefaultDOP: + description: "Default degree of parallelism. This an internal experimental parameter and should not be changed on live systems." + set_at: [ startup, runtime ] + cpp_varname: "internalQueryDefaultDOP" + cpp_vartype: AtomicWord<int> + default: 1 + test_only: true + validator: + gt: 0 + diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp index bd75a303166..f97bb74f065 100644 --- a/src/mongo/db/query/query_planner.cpp +++ b/src/mongo/db/query/query_planner.cpp @@ -57,6 +57,47 @@ #include "mongo/logv2/log.h" namespace mongo { +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 using std::numeric_limits; using std::unique_ptr; @@ -1102,4 +1143,195 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan( return {std::move(out)}; } +StatusWith<QueryPlanner::SubqueriesPlanningResult> QueryPlanner::planSubqueries( + OperationContext* opCtx, + const Collection* collection, + const PlanCache* planCache, + const CanonicalQuery& query, + const QueryPlannerParams& params) { + invariant(query.root()->matchType() == MatchExpression::OR); + invariant(query.root()->numChildren(), "Cannot plan subqueries for an $or with no children"); + + SubqueriesPlanningResult planningResult{query.root()->shallowClone()}; + for (size_t i = 0; i < params.indices.size(); ++i) { + const IndexEntry& ie = params.indices[i]; + const auto insertionRes = planningResult.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 < planningResult.orExpression->numChildren(); ++i) { + // We need a place to shove the results from planning this branch. + planningResult.branches.push_back( + std::make_unique<SubqueriesPlanningResult::BranchPlanningResult>()); + auto branchResult = planningResult.branches.back().get(); + auto orChild = planningResult.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. + + // Populate branchResult->cachedSolution if an active cachedSolution entry exists. + if (planCache && 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 = planningResult.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 = + planningResult.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, params); + 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 std::move(planningResult); +} + +StatusWith<std::unique_ptr<QuerySolution>> QueryPlanner::choosePlanForSubqueries( + const CanonicalQuery& query, + const QueryPlannerParams& params, + QueryPlanner::SubqueriesPlanningResult planningResult, + std::function<StatusWith<std::unique_ptr<QuerySolution>>( + CanonicalQuery* cq, std::vector<unique_ptr<QuerySolution>>)> multiplanCallback) { + // 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 < planningResult.orExpression->numChildren(); ++i) { + auto orChild = planningResult.orExpression->getChild(i); + auto branchResult = planningResult.branches[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, + planningResult.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, planningResult.indexMap); + if (!tagStatus.isOK()) { + return tagStatus; + } + } else { + // N solutions, rank them. + + invariant(!branchResult->solutions.empty()); + + auto multiPlanStatus = multiplanCallback(branchResult->canonicalQuery.get(), + std::move(branchResult->solutions)); + if (!multiPlanStatus.isOK()) { + return multiPlanStatus; + } + + auto bestSoln = std::move(multiPlanStatus.getValue()); + + // 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(), planningResult.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(planningResult.orExpression.get()); + + // Use the cached index assignments to build solnRoot. Takes ownership of '_orExpression'. + std::unique_ptr<QuerySolutionNode> solnRoot(QueryPlannerAccess::buildIndexedDataAccess( + query, std::move(planningResult.orExpression), params.indices, params)); + + 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())); + + auto compositeSolution = + QueryPlannerAnalysis::analyzeDataAccess(query, params, 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())); + + return std::move(compositeSolution); +} } // namespace mongo diff --git a/src/mongo/db/query/query_planner.h b/src/mongo/db/query/query_planner.h index f57fb876ea9..e9d27bea0f5 100644 --- a/src/mongo/db/query/query_planner.h +++ b/src/mongo/db/query/query_planner.h @@ -45,6 +45,40 @@ class Collection; */ class QueryPlanner { public: + /** + * Holds the result of subqueries planning for rooted $or queries. + */ + struct SubqueriesPlanningResult { + /** + * A class used internally in order to keep track of the results of planning + * a particular $or branch. + */ + struct 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; + }; + + // 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; + + // Holds a list of the results from planning each branch. + std::vector<std::unique_ptr<BranchPlanningResult>> branches; + + // We need this to extract cache-friendly index data from the index assignments. + std::map<IndexEntry::Identifier, size_t> indexMap; + }; + // Identifies the version of the query planner module. Reported in explain. static const int kPlannerVersion; @@ -68,6 +102,16 @@ public: const CachedSolution& cachedSoln); /** + * Plan each branch of the rooted $or query independently, and store the resulting + * lists of query solutions in 'SubqueriesPlanningResult'. + */ + static StatusWith<SubqueriesPlanningResult> planSubqueries(OperationContext* opCtx, + const Collection* collection, + const PlanCache* planCache, + const CanonicalQuery& query, + const QueryPlannerParams& params); + + /** * Generates and returns the index tag tree that will be inserted into the plan cache. This data * gets stashed inside a QuerySolution until it can be inserted into the cache proper. * @@ -99,6 +143,18 @@ public: static Status tagAccordingToCache(MatchExpression* filter, const PlanCacheIndexTree* const indexTree, const std::map<IndexEntry::Identifier, size_t>& indexMap); -}; + /** + * Uses the query planning results from QueryPlanner::planSubqueries() and the multi planner + * callback to select the best plan for each branch. + * + * On success, returns a composite solution obtained by planning each $or branch independently. + */ + static StatusWith<std::unique_ptr<QuerySolution>> choosePlanForSubqueries( + const CanonicalQuery& query, + const QueryPlannerParams& params, + QueryPlanner::SubqueriesPlanningResult planningResult, + std::function<StatusWith<std::unique_ptr<QuerySolution>>( + CanonicalQuery* cq, std::vector<std::unique_ptr<QuerySolution>>)> multiplanCallback); +}; } // namespace mongo diff --git a/src/mongo/db/query/query_planner_common.cpp b/src/mongo/db/query/query_planner_common.cpp index 55b8b0e71c4..bba5b5d6eb5 100644 --- a/src/mongo/db/query/query_planner_common.cpp +++ b/src/mongo/db/query/query_planner_common.cpp @@ -33,12 +33,11 @@ #include "mongo/base/exact_cast.h" #include "mongo/db/query/projection_ast_path_tracking_visitor.h" -#include "mongo/db/query/projection_ast_walker.h" #include "mongo/db/query/query_planner_common.h" +#include "mongo/db/query/tree_walker.h" #include "mongo/logv2/redaction.h" #include "mongo/util/assert_util.h" - namespace mongo { void QueryPlannerCommon::reverseScans(QuerySolutionNode* node) { @@ -123,7 +122,7 @@ std::vector<FieldPath> QueryPlannerCommon::extractSortKeyMetaFieldsFromProjectio MetaFieldVisitorContext ctx; MetaFieldVisitor visitor(&ctx); projection_ast::PathTrackingConstWalker<MetaFieldData> walker{&ctx, {&visitor}, {}}; - projection_ast_walker::walk(&walker, proj.root()); + tree_walker::walk<true, projection_ast::ASTNode>(proj.root(), &walker); return std::move(ctx.data().metaPaths); } diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp index 1d51d66022b..0e62859febf 100644 --- a/src/mongo/db/query/query_solution.cpp +++ b/src/mongo/db/query/query_solution.cpp @@ -134,6 +134,20 @@ void QuerySolutionNode::addCommon(str::stream* ss, int indent) const { *ss << "providedSorts = {" << providedSorts().debugString() << "}" << '\n'; } +bool QuerySolutionNode::hasNode(StageType type) const { + if (type == getType()) { + return true; + } + + for (auto&& child : children) { + if (child->hasNode(type)) { + return true; + } + } + + return false; +} + // // TextNode // diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h index 4d8e5c06ee7..cdb038cd61e 100644 --- a/src/mongo/db/query/query_solution.h +++ b/src/mongo/db/query/query_solution.h @@ -233,6 +233,11 @@ struct QuerySolutionNode { [](auto& child) { return child.release(); }); } + /** + * True, if this node, or any of it's children is of the given 'type'. + */ + bool hasNode(StageType type) const; + // These are owned here. // // TODO SERVER-35512: Make this a vector of unique_ptr. @@ -314,6 +319,13 @@ struct QuerySolution { std::unique_ptr<SolutionCacheData> cacheData; /** + * True, of this solution tree contains a node of the given 'type'. + */ + bool hasNode(StageType type) const { + return root && root->hasNode(type); + } + + /** * Output a human-readable std::string representing the plan. */ std::string toString() { diff --git a/src/mongo/db/query/sbe_cached_solution_planner.cpp b/src/mongo/db/query/sbe_cached_solution_planner.cpp new file mode 100644 index 00000000000..bf8a6613463 --- /dev/null +++ b/src/mongo/db/query/sbe_cached_solution_planner.cpp @@ -0,0 +1,158 @@ +/** + * 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/query/sbe_cached_solution_planner.h" + +#include "mongo/db/query/collection_query_info.h" +#include "mongo/db/query/explain.h" +#include "mongo/db/query/query_planner.h" +#include "mongo/db/query/sbe_multi_planner.h" +#include "mongo/db/query/stage_builder_util.h" +#include "mongo/logv2/log.h" + +namespace mongo::sbe { +plan_ranker::CandidatePlan CachedSolutionPlanner::plan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) { + invariant(solutions.size() == 1); + invariant(solutions.size() == roots.size()); + + auto candidate = [&]() { + auto candidates = collectExecutionStats(std::move(solutions), std::move(roots)); + invariant(candidates.size() == 1); + return std::move(candidates[0]); + }(); + + if (candidate.failed) { + // On failure, fall back to replanning the whole query. We neither evict the existing cache + // entry, nor cache the result of replanning. + LOGV2_DEBUG(2057901, + 1, + "Execution of cached plan failed, falling back to replan. query: " + "{canonicalQuery_Short} planSummary: {Explain_getPlanSummary_child_get} ", + "canonicalQuery_Short"_attr = redact(_cq.toStringShort()), + "Explain_getPlanSummary_child_get"_attr = + Explain::getPlanSummary(candidate.root.get())); + return replan(false); + } + + auto stats{candidate.root->getStats()}; + auto numReads{calculateNumberOfReads(stats.get())}; + // If the cached plan hit EOF quickly enough, or still as efficient as before, then no need to + // replan. Finalize the cached plan and return it. + if (stats->common.isEOF || numReads <= _decisionReads) { + return finalizeExecutionPlan(std::move(stats), std::move(candidate)); + } + + // If we're here, the trial period took more than 'maxReadsBeforeReplan' physical reads. This + // plan may not be efficient any longer, so we replan from scratch. + LOGV2_DEBUG( + 2058001, + 1, + "Execution of cached plan required {maxReadsBeforeReplan} works, but was originally cached " + "with only {decisionReads} works. Evicting cache entry and replanning query: " + "{canonicalQuery_Short} plan summary before replan: {Explain_getPlanSummary_child_get}", + "maxReadsBeforeReplan"_attr = numReads, + "decisionReads"_attr = _decisionReads, + "canonicalQuery_Short"_attr = redact(_cq.toStringShort()), + "Explain_getPlanSummary_child_get"_attr = Explain::getPlanSummary(candidate.root.get())); + return replan(true); +} + +plan_ranker::CandidatePlan CachedSolutionPlanner::finalizeExecutionPlan( + std::unique_ptr<PlanStageStats> stats, plan_ranker::CandidatePlan candidate) const { + // If the winning stage has exited early, clear the results queue and reopen the plan stage + // tree, as we cannot resume such execution tree from where the trial run has stopped, and, as + // a result, we cannot stash the results returned so far in the plan executor. + if (!stats->common.isEOF && candidate.exitedEarly) { + candidate.root->close(); + candidate.root->open(true); + // Clear the results queue. + candidate.results = decltype(candidate.results){}; + } + + return candidate; +} + +plan_ranker::CandidatePlan CachedSolutionPlanner::replan(bool shouldCache) const { + if (shouldCache) { + // Deactivate the current cache entry. + auto cache = CollectionQueryInfo::get(_collection).getPlanCache(); + cache->deactivate(_cq); + } + + // Use the query planning module to plan the whole query. + auto solutions = uassertStatusOK(QueryPlanner::plan(_cq, _queryParams)); + if (solutions.size() == 1) { + // Only one possible plan. Build the stages from the solution. + auto&& [root, data] = stage_builder::buildSlotBasedExecutableTree( + _opCtx, _collection, _cq, *solutions[0], _yieldPolicy, true); + prepareExecutionPlan(root.get(), &data); + LOGV2_DEBUG( + 2058101, + 1, + "Replanning of query resulted in single query solution, which will not be cached. " + "{canonicalQuery_Short} plan summary after replan: {Explain_getPlanSummary_child_get} " + "previous cache entry evicted: {shouldCache_yes_no}", + "canonicalQuery_Short"_attr = redact(_cq.toStringShort()), + "Explain_getPlanSummary_child_get"_attr = Explain::getPlanSummary(root.get()), + "shouldCache_yes_no"_attr = (shouldCache ? "yes" : "no")); + return {std::move(solutions[0]), std::move(root), std::move(data)}; + } + + // Many solutions. Build a plan stage tree for each solution and create a multi planner to pick + // the best, update the cache, and so on. + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots; + for (auto&& solution : solutions) { + if (solution->cacheData.get()) { + solution->cacheData->indexFilterApplied = _queryParams.indexFiltersApplied; + } + + roots.push_back(stage_builder::buildSlotBasedExecutableTree( + _opCtx, _collection, _cq, *solution, _yieldPolicy, true)); + } + + const auto cachingMode = + shouldCache ? PlanCachingMode::AlwaysCache : PlanCachingMode::NeverCache; + MultiPlanner multiPlanner{_opCtx, _collection, _cq, cachingMode, _yieldPolicy}; + auto plan = multiPlanner.plan(std::move(solutions), std::move(roots)); + LOGV2_DEBUG(2058201, + 1, + "Replanning {canonicalQuery_Short} resulted in plan with summary: " + "{Explain_getPlanSummary_child_get}, which {shouldCache_has_has_not} been written " + "to the cache", + "canonicalQuery_Short"_attr = redact(_cq.toStringShort()), + "Explain_getPlanSummary_child_get"_attr = Explain::getPlanSummary(plan.root.get()), + "shouldCache_has_has_not"_attr = (shouldCache ? "has" : "has not")); + return plan; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/query/sbe_cached_solution_planner.h b/src/mongo/db/query/sbe_cached_solution_planner.h new file mode 100644 index 00000000000..d14c44e2d72 --- /dev/null +++ b/src/mongo/db/query/sbe_cached_solution_planner.h @@ -0,0 +1,86 @@ +/** + * 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/sbe_plan_ranker.h" +#include "mongo/db/query/sbe_runtime_planner.h" + +namespace mongo::sbe { +/** + * Runs a trial period in order to evaluate the cost of a cached plan. If the cost is unexpectedly + * high, the plan cache entry is deactivated and we use multi-planning to select an entirely new + * winning plan. This process is called "replanning". + * + * TODO: refresh the list of indexes in 'queryParams' during replanning. + */ +class CachedSolutionPlanner final : public BaseRuntimePlanner { +public: + CachedSolutionPlanner(OperationContext* opCtx, + Collection* collection, + const CanonicalQuery& cq, + const QueryPlannerParams& queryParams, + size_t decisionReads, + PlanYieldPolicySBE* yieldPolicy) + : BaseRuntimePlanner{opCtx, collection, cq, yieldPolicy}, + _queryParams{queryParams}, + _decisionReads{decisionReads} {} + + plan_ranker::CandidatePlan plan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) + final; + +private: + /** + * Finalizes the winning plan before passing it to the caller as a result of the planning. + */ + plan_ranker::CandidatePlan finalizeExecutionPlan(std::unique_ptr<sbe::PlanStageStats> stats, + plan_ranker::CandidatePlan candidate) const; + + /** + * Uses the QueryPlanner and the MultiPlanner to re-generate candidate plans for this + * query and select a new winner. + * + * Falls back to a new plan if the performance was worse than anticipated during the trial + * period. + * + * The plan cache is modified only if 'shouldCache' is true. + */ + plan_ranker::CandidatePlan replan(bool shouldCache) const; + + // Query parameters used to create a query solution when the plan was first created. Used during + // replanning. + const QueryPlannerParams _queryParams; + + // The number of physical reads taken to decide on a winning plan when the plan was first + // cached. + const size_t _decisionReads; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/query/sbe_multi_planner.cpp b/src/mongo/db/query/sbe_multi_planner.cpp new file mode 100644 index 00000000000..e6d4aedb39f --- /dev/null +++ b/src/mongo/db/query/sbe_multi_planner.cpp @@ -0,0 +1,95 @@ +/** + * 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/query/sbe_multi_planner.h" + +#include "mongo/db/exec/multi_plan.h" +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/query/collection_query_info.h" +#include "mongo/db/query/explain.h" +#include "mongo/db/query/query_planner.h" +#include "mongo/db/query/stage_builder_util.h" +#include "mongo/logv2/log.h" + +namespace mongo::sbe { +plan_ranker::CandidatePlan MultiPlanner::plan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) { + auto candidates = collectExecutionStats(std::move(solutions), std::move(roots)); + auto decision = uassertStatusOK(mongo::plan_ranker::pickBestPlan<PlanStageStats>(candidates)); + return finalizeExecutionPlans(std::move(decision), std::move(candidates)); +} + +plan_ranker::CandidatePlan MultiPlanner::finalizeExecutionPlans( + std::unique_ptr<mongo::plan_ranker::PlanRankingDecision> decision, + std::vector<plan_ranker::CandidatePlan> candidates) const { + invariant(decision); + + auto&& stats = decision->getStats<sbe::PlanStageStats>(); + const auto winnerIdx = decision->candidateOrder[0]; + invariant(winnerIdx < candidates.size()); + invariant(winnerIdx < stats.size()); + auto& winner = candidates[winnerIdx]; + + LOGV2_DEBUG( + 4822875, 5, "Winning solution", "bestSolution"_attr = redact(winner.solution->toString())); + LOGV2_DEBUG(4822876, + 2, + "Winning plan", + "planSumamry"_attr = Explain::getPlanSummary(winner.root.get())); + + // Close all candidate plans but the winner. + for (size_t ix = 1; ix < decision->candidateOrder.size(); ++ix) { + const auto planIdx = decision->candidateOrder[ix]; + invariant(planIdx < candidates.size()); + candidates[planIdx].root->close(); + } + + // If the winning stage has exited early but has not fetched all results, clear the results + // queue and reopen the plan stage tree, as we cannot resume such execution tree from where + // the trial run has stopped, and, as a result, we cannot stash the results returned so far + // in the plan executor. + if (!stats[winnerIdx]->common.isEOF && winner.exitedEarly) { + winner.root->close(); + winner.root->open(true); + // Clear the results queue. + winner.results = decltype(winner.results){}; + } + + // Writes a cache entry for the winning plan to the plan cache if possible. + plan_cache_util::updatePlanCache( + _opCtx, _collection, _cachingMode, _cq, std::move(decision), candidates); + + return std::move(winner); +} +} // namespace mongo::sbe diff --git a/src/mongo/db/query/sbe_multi_planner.h b/src/mongo/db/query/sbe_multi_planner.h new file mode 100644 index 00000000000..48dca7081bf --- /dev/null +++ b/src/mongo/db/query/sbe_multi_planner.h @@ -0,0 +1,71 @@ +/** + * 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_cache_util.h" +#include "mongo/db/query/plan_executor.h" +#include "mongo/db/query/plan_ranker.h" +#include "mongo/db/query/sbe_runtime_planner.h" + +namespace mongo::sbe { +/** + * Collects execution stats for all candidate plans, ranks them and picks the best. + * + * TODO: add support for backup plan + */ +class MultiPlanner final : public BaseRuntimePlanner { +public: + MultiPlanner(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + PlanCachingMode cachingMode, + PlanYieldPolicySBE* yieldPolicy) + : BaseRuntimePlanner{opCtx, collection, cq, yieldPolicy}, _cachingMode{cachingMode} {} + + plan_ranker::CandidatePlan plan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) + final; + +private: + /** + * Returns the best candidate plan selected according to the plan ranking 'decision'. + * + * Calls 'close' method on all other candidate plans and updates the plan cache entry, + * if possible. + */ + plan_ranker::CandidatePlan finalizeExecutionPlans( + std::unique_ptr<mongo::plan_ranker::PlanRankingDecision> decision, + std::vector<plan_ranker::CandidatePlan> candidates) const; + + // Describes the cases in which we should write an entry for the winning plan to the plan cache. + const PlanCachingMode _cachingMode; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/query/sbe_plan_ranker.cpp b/src/mongo/db/query/sbe_plan_ranker.cpp new file mode 100644 index 00000000000..478e9e64360 --- /dev/null +++ b/src/mongo/db/query/sbe_plan_ranker.cpp @@ -0,0 +1,93 @@ +/** + * 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/query/sbe_plan_ranker.h" + +namespace mongo::sbe::plan_ranker { +namespace { +/** + * A plan ranker for the SBE plan stage tree. Defines productivity as a cumulative number of + * physical reads from the storage performed by all stages in the plan which can read from the + * storage, divided by the total number of advances of the root stage, which corresponds to the + * number of returned documents. + */ +class DefaultPlanScorer final : public mongo::plan_ranker::PlanScorer<PlanStageStats> { +public: + DefaultPlanScorer(const QuerySolution* solution) : _solution{solution} { + invariant(_solution); + } + +protected: + double calculateProductivity(const mongo::sbe::PlanStageStats* root) const final { + auto numReads{calculateNumberOfReads(root)}; + + if (numReads == 0) { + return 0; + } + + return static_cast<double>(root->common.advances) / static_cast<double>(numReads); + } + + std::string getProductivityFormula(const mongo::sbe::PlanStageStats* root) const final { + auto numReads{calculateNumberOfReads(root)}; + StringBuilder sb; + + if (numReads == 0) { + sb << "(0 numReads)"; + } else { + sb << "(" << root->common.advances << " advances)/(" << numReads << " numReads)"; + } + + return sb.str(); + } + + double getNumberOfAdvances(const mongo::sbe::PlanStageStats* stats) const final { + return stats->common.advances; + } + + bool hasStage(StageType type, const mongo::sbe::PlanStageStats* stats) const final { + // In SBE a plan stage doesn't map 1-to-1 to a solution node, and can expand into a subtree + // of plan stages, each having its own plan stage stats. So, to answer whether an SBE plan + // stage stats tree contains a stage of the given 'type', we need to look into the solution + // tree instead. + return _solution->hasNode(type); + } + +private: + const QuerySolution* _solution; +}; +} // namespace + +std::unique_ptr<mongo::plan_ranker::PlanScorer<PlanStageStats>> makePlanScorer( + const QuerySolution* solution) { + return std::make_unique<DefaultPlanScorer>(solution); +} +} // namespace mongo::sbe::plan_ranker diff --git a/src/mongo/db/query/sbe_plan_ranker.h b/src/mongo/db/query/sbe_plan_ranker.h new file mode 100644 index 00000000000..9a08a2888fa --- /dev/null +++ b/src/mongo/db/query/sbe_plan_ranker.h @@ -0,0 +1,47 @@ +/** + * 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/plan_ranker.h" +#include "mongo/db/query/sbe_stage_builder.h" + +namespace mongo::sbe::plan_ranker { + +using CandidatePlan = + mongo::plan_ranker::BaseCandidatePlan<std::unique_ptr<mongo::sbe::PlanStage>, + std::pair<BSONObj, boost::optional<RecordId>>, + stage_builder::PlanStageData>; + +/** + * A factory function to create a plan ranker for an SBE plan stage stats tree. + */ +std::unique_ptr<mongo::plan_ranker::PlanScorer<PlanStageStats>> makePlanScorer( + const QuerySolution* solution); +} // namespace mongo::sbe::plan_ranker diff --git a/src/mongo/db/query/sbe_runtime_planner.cpp b/src/mongo/db/query/sbe_runtime_planner.cpp new file mode 100644 index 00000000000..254cfeb84de --- /dev/null +++ b/src/mongo/db/query/sbe_runtime_planner.cpp @@ -0,0 +1,157 @@ +/** + * 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/query/sbe_runtime_planner.h" + +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/trial_period_utils.h" +#include "mongo/db/query/plan_executor_sbe.h" + +namespace mongo::sbe { +namespace { +/** + * Fetches a next document form the given plan stage tree and returns 'true' if the plan stage + * returns EOF, or throws 'TrialRunProgressTracker::EarlyExitException' exception. Otherwise, the + * loaded document is placed into the candidate's plan result queue. + * + * If the plan stage throws a 'QueryExceededMemoryLimitNoDiskUseAllowed', it will be caught and the + * 'candidate->failed' flag will be set to 'true', and the 'numFailures' parameter incremented by 1. + * This failure is considered recoverable, as another candidate plan may require less memory, or may + * not contain a stage requiring spilling to disk at all. + */ +bool fetchNextDocument(plan_ranker::CandidatePlan* candidate, + const std::pair<sbe::value::SlotAccessor*, sbe::value::SlotAccessor*>& slots, + size_t* numFailures) { + try { + BSONObj obj; + RecordId recordId; + + auto [resultSlot, recordIdSlot] = slots; + auto state = fetchNext(candidate->root.get(), + resultSlot, + recordIdSlot, + &obj, + recordIdSlot ? &recordId : nullptr); + if (state == sbe::PlanState::IS_EOF) { + candidate->root->close(); + return true; + } + + invariant(state == sbe::PlanState::ADVANCED); + candidate->results.push({std::move(obj), {recordIdSlot != nullptr, recordId}}); + } catch (const ExceptionFor<ErrorCodes::QueryTrialRunCompleted>&) { + candidate->exitedEarly = true; + return true; + } catch (const ExceptionFor<ErrorCodes::QueryExceededMemoryLimitNoDiskUseAllowed>&) { + candidate->root->close(); + candidate->failed = true; + ++(*numFailures); + } + return false; +} +} // namespace + +std::tuple<sbe::value::SlotAccessor*, sbe::value::SlotAccessor*, bool> +BaseRuntimePlanner::prepareExecutionPlan(PlanStage* root, + stage_builder::PlanStageData* data) const { + invariant(root); + invariant(data); + + root->prepare(data->ctx); + + sbe::value::SlotAccessor* resultSlot{nullptr}; + if (data->resultSlot) { + resultSlot = root->getAccessor(data->ctx, *data->resultSlot); + uassert(4822871, "Query does not have result slot.", resultSlot); + } + + sbe::value::SlotAccessor* recordIdSlot{nullptr}; + if (data->recordIdSlot) { + recordIdSlot = root->getAccessor(data->ctx, *data->recordIdSlot); + uassert(4822872, "Query does not have record ID slot.", recordIdSlot); + } + + root->attachFromOperationContext(_opCtx); + + auto exitedEarly{false}; + try { + root->open(false); + } catch (const ExceptionFor<ErrorCodes::QueryTrialRunCompleted>&) { + exitedEarly = true; + } + + return {resultSlot, recordIdSlot, exitedEarly}; +} + +std::vector<plan_ranker::CandidatePlan> BaseRuntimePlanner::collectExecutionStats( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) { + invariant(solutions.size() == roots.size()); + + std::vector<plan_ranker::CandidatePlan> candidates; + std::vector<std::pair<sbe::value::SlotAccessor*, sbe::value::SlotAccessor*>> slots; + + for (size_t ix = 0; ix < roots.size(); ++ix) { + auto&& [root, data] = roots[ix]; + auto [resultSlot, recordIdSlot, exitedEarly] = prepareExecutionPlan(root.get(), &data); + + candidates.push_back( + {std::move(solutions[ix]), std::move(root), std::move(data), exitedEarly}); + slots.push_back({resultSlot, recordIdSlot}); + } + + auto done{false}; + size_t numFailures{0}; + const auto maxNumResults{trial_period::getTrialPeriodNumToReturn(_cq)}; + for (size_t it = 0; it < maxNumResults && !done; ++it) { + for (size_t ix = 0; ix < candidates.size(); ++ix) { + // Even if we had a candidate plan that exited early, we still want continue the trial + // run as the early exited plan may not be the best. E.g., it could be blocked in a SORT + // stage until one of the trial period metrics was reached, causing the plan to raise an + // early exit exception and return control back to the runtime planner. If that happens, + // we need to continue and complete the trial period for all candidates, as some of them + // may have a better cost. + if (candidates[ix].failed || candidates[ix].exitedEarly) { + continue; + } + + done |= fetchNextDocument(&candidates[ix], slots[ix], &numFailures) || + (numFailures == candidates.size()); + } + } + + // Make sure we have at least one plan which hasn't failed. + uassert(4822873, + "Runtime planner encountered a failure while collecting execution stats", + numFailures != candidates.size()); + + return candidates; +} +} // namespace mongo::sbe diff --git a/src/mongo/db/query/sbe_runtime_planner.h b/src/mongo/db/query/sbe_runtime_planner.h new file mode 100644 index 00000000000..d399b096877 --- /dev/null +++ b/src/mongo/db/query/sbe_runtime_planner.h @@ -0,0 +1,102 @@ +/** + * 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/canonical_query.h" +#include "mongo/db/query/plan_yield_policy_sbe.h" +#include "mongo/db/query/query_solution.h" +#include "mongo/db/query/sbe_plan_ranker.h" + +namespace mongo::sbe { +/** + * An interface to be implemented by all classes which can evaluate the cost of a PlanStage tree in + * order to pick the the best plan amongst those specified in 'roots' vector. Evaluation is done in + * runtime by collecting execution stats for each of the plans, and the best candidate plan is + * chosen according to certain criteria. + */ +class RuntimePlanner { +public: + virtual ~RuntimePlanner() = default; + + virtual plan_ranker::CandidatePlan plan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) = 0; +}; + +/** + * A base class for runtime planner which provides a method to perform a trial run for the candidate + * plan by executing each plan in a round-robin fashion and collecting execution stats. Each + * specific implementation can use the collected stats to select the best plan amongst the + * candidates. + */ +class BaseRuntimePlanner : public RuntimePlanner { +public: + BaseRuntimePlanner(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + PlanYieldPolicySBE* yieldPolicy) + : _opCtx(opCtx), _collection(collection), _cq(cq), _yieldPolicy(yieldPolicy) { + invariant(_opCtx); + invariant(_collection); + } + +protected: + /** + * Prepares the given plan stage tree for execution, attaches it to the operation context and + * returns two slot accessors for the result and recordId slots, and a boolean value indicating + * if the plan has exited early from the trial period. + */ + std::tuple<sbe::value::SlotAccessor*, sbe::value::SlotAccessor*, bool> prepareExecutionPlan( + PlanStage* root, stage_builder::PlanStageData* data) const; + + /** + * Executes each plan in a round-robin fashion to collect execution stats. Stops when: + * * Any plan hits EOF. + * * Or returns a pre-defined number of results. + * * Or all candidate plans fail or exit early by throwing a special signaling exception. + * + * All documents returned by each plan are enqueued into the 'CandidatePlan->results' queue. + * + * Upon completion returns a vector of candidate plans. Execution stats can be obtained for each + * of the candidate plans by calling 'CandidatePlan->root->getStats()'. + * + * After the trial period ends, all plans remain open. + */ + std::vector<plan_ranker::CandidatePlan> collectExecutionStats( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots); + + OperationContext* const _opCtx; + const Collection* const _collection; + const CanonicalQuery& _cq; + PlanYieldPolicySBE* const _yieldPolicy; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp new file mode 100644 index 00000000000..fcacfaa402e --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder.cpp @@ -0,0 +1,440 @@ +/** + * 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/query/sbe_stage_builder.h" + +#include "mongo/db/catalog/collection.h" +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/hash_agg.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/db/exec/sbe/stages/makeobj.h" +#include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/stages/scan.h" +#include "mongo/db/exec/sbe/stages/sort.h" +#include "mongo/db/exec/sbe/stages/text_match.h" +#include "mongo/db/exec/sbe/stages/traverse.h" +#include "mongo/db/exec/sbe/stages/union.h" +#include "mongo/db/fts/fts_index_format.h" +#include "mongo/db/fts/fts_query_impl.h" +#include "mongo/db/fts/fts_spec.h" +#include "mongo/db/index/fts_access_method.h" +#include "mongo/db/query/sbe_stage_builder_coll_scan.h" +#include "mongo/db/query/sbe_stage_builder_filter.h" +#include "mongo/db/query/sbe_stage_builder_index_scan.h" +#include "mongo/db/query/sbe_stage_builder_projection.h" + +namespace mongo::stage_builder { +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildCollScan( + const QuerySolutionNode* root) { + auto csn = static_cast<const CollectionScanNode*>(root); + auto [resultSlot, recordIdSlot, oplogTsSlot, stage] = + generateCollScan(_opCtx, + _collection, + csn, + &_slotIdGenerator, + _yieldPolicy, + _data.trialRunProgressTracker.get()); + _data.resultSlot = resultSlot; + _data.recordIdSlot = recordIdSlot; + _data.oplogTsSlot = oplogTsSlot; + _data.shouldTrackLatestOplogTimestamp = csn->shouldTrackLatestOplogTimestamp; + _data.shouldTrackResumeToken = csn->requestResumeToken; + return std::move(stage); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildIndexScan( + const QuerySolutionNode* root) { + auto ixn = static_cast<const IndexScanNode*>(root); + auto [slot, stage] = generateIndexScan(_opCtx, + _collection, + ixn, + &_slotIdGenerator, + &_spoolIdGenerator, + _yieldPolicy, + _data.trialRunProgressTracker.get()); + _data.recordIdSlot = slot; + return std::move(stage); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::makeLoopJoinForFetch( + std::unique_ptr<sbe::PlanStage> inputStage, sbe::value::SlotId recordIdKeySlot) { + _data.resultSlot = _slotIdGenerator.generate(); + _data.recordIdSlot = _slotIdGenerator.generate(); + + // Scan the collection in the range [recordIdKeySlot, Inf). + auto scanStage = sbe::makeS<sbe::ScanStage>( + NamespaceStringOrUUID{_collection->ns().db().toString(), _collection->uuid()}, + _data.resultSlot, + _data.recordIdSlot, + std::vector<std::string>{}, + sbe::makeSV(), + recordIdKeySlot, + true, + nullptr, + _data.trialRunProgressTracker.get()); + + // Get the recordIdKeySlot from the outer side (e.g., IXSCAN) and feed it to the inner side, + // limiting the result set to 1 row. + return sbe::makeS<sbe::LoopJoinStage>( + std::move(inputStage), + sbe::makeS<sbe::LimitSkipStage>(std::move(scanStage), 1, boost::none), + sbe::makeSV(), + sbe::makeSV(recordIdKeySlot), + nullptr); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildFetch(const QuerySolutionNode* root) { + auto fn = static_cast<const FetchNode*>(root); + auto inputStage = build(fn->children[0]); + + uassert(4822880, "RecordId slot is not defined", _data.recordIdSlot); + + auto stage = makeLoopJoinForFetch(std::move(inputStage), *_data.recordIdSlot); + + if (fn->filter) { + stage = generateFilter( + fn->filter.get(), std::move(stage), &_slotIdGenerator, *_data.resultSlot); + } + + return stage; +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildLimit(const QuerySolutionNode* root) { + const auto ln = static_cast<const LimitNode*>(root); + // If we have both limit and skip stages and the skip stage is beneath the limit, then we can + // combine these two stages into one. So, save the _limit value and let the skip stage builder + // handle it. + if (ln->children[0]->getType() == StageType::STAGE_SKIP) { + _limit = ln->limit; + } + auto inputStage = build(ln->children[0]); + return _limit + ? std::move(inputStage) + : std::make_unique<sbe::LimitSkipStage>(std::move(inputStage), ln->limit, boost::none); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildSkip(const QuerySolutionNode* root) { + const auto sn = static_cast<const SkipNode*>(root); + auto inputStage = build(sn->children[0]); + return std::make_unique<sbe::LimitSkipStage>(std::move(inputStage), _limit, sn->skip); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildSort(const QuerySolutionNode* root) { + // TODO SERVER-48470: Replace std::string_view with StringData. + using namespace std::literals; + + const auto sn = static_cast<const SortNode*>(root); + auto sortPattern = SortPattern{sn->pattern, _cq.getExpCtx()}; + auto inputStage = build(sn->children[0]); + sbe::value::SlotVector orderBy; + std::vector<sbe::value::SortDirection> direction; + sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projectMap; + + for (const auto& part : sortPattern) { + uassert(4822881, "Sorting by expression not supported", !part.expression); + uassert(4822882, + "Sorting by dotted paths not supported", + part.fieldPath && part.fieldPath->getPathLength() == 1); + + // Slot holding the sort key. + auto sortFieldVar{_slotIdGenerator.generate()}; + orderBy.push_back(sortFieldVar); + direction.push_back(part.isAscending ? sbe::value::SortDirection::Ascending + : sbe::value::SortDirection::Descending); + + // Generate projection to get the value of the sort key. Ideally, this should be + // tracked by a 'reference tracker' at higher level. + auto fieldName = part.fieldPath->getFieldName(0); + auto fieldNameSV = std::string_view{fieldName.rawData(), fieldName.size()}; + projectMap.emplace( + sortFieldVar, + sbe::makeE<sbe::EFunction>("getField"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(*_data.resultSlot), + sbe::makeE<sbe::EConstant>(fieldNameSV)))); + } + + inputStage = sbe::makeS<sbe::ProjectStage>(std::move(inputStage), std::move(projectMap)); + + // Generate traversals to pick the min/max element from arrays. + for (size_t idx = 0; idx < orderBy.size(); ++idx) { + auto resultVar{_slotIdGenerator.generate()}; + auto innerVar{_slotIdGenerator.generate()}; + + auto innerBranch = sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + innerVar, + sbe::makeE<sbe::EVariable>(orderBy[idx])); + + auto op = direction[idx] == sbe::value::SortDirection::Ascending + ? sbe::EPrimBinary::less + : sbe::EPrimBinary::greater; + auto minmax = sbe::makeE<sbe::EIf>( + sbe::makeE<sbe::EPrimBinary>( + op, + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::cmp3w, + sbe::makeE<sbe::EVariable>(innerVar), + sbe::makeE<sbe::EVariable>(resultVar)), + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt64, 0)), + sbe::makeE<sbe::EVariable>(innerVar), + sbe::makeE<sbe::EVariable>(resultVar)); + + inputStage = sbe::makeS<sbe::TraverseStage>(std::move(inputStage), + std::move(innerBranch), + orderBy[idx], + resultVar, + innerVar, + sbe::makeSV(), + std::move(minmax), + nullptr); + orderBy[idx] = resultVar; + } + + sbe::value::SlotVector values; + values.push_back(*_data.resultSlot); + if (_data.recordIdSlot) { + // Break ties with record id if available. + orderBy.push_back(*_data.recordIdSlot); + // This is arbitrary. + direction.push_back(sbe::value::SortDirection::Ascending); + } + + // A sort stage is a binding reflector, so we need to plumb through the 'oplogTsSlot' to make + // it visible at the root stage. + if (_data.oplogTsSlot) { + values.push_back(*_data.oplogTsSlot); + } + + return sbe::makeS<sbe::SortStage>(std::move(inputStage), + std::move(orderBy), + std::move(direction), + std::move(values), + sn->limit ? sn->limit + : std::numeric_limits<std::size_t>::max(), + _data.trialRunProgressTracker.get()); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildSortKeyGeneraror( + const QuerySolutionNode* root) { + uasserted(4822883, "Sort key generator in not supported in SBE yet"); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildProjectionSimple( + const QuerySolutionNode* root) { + using namespace std::literals; + + auto pn = static_cast<const ProjectionNodeSimple*>(root); + auto inputStage = build(pn->children[0]); + sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projections; + sbe::value::SlotVector fieldSlots; + + for (const auto& field : pn->proj.getRequiredFields()) { + fieldSlots.push_back(_slotIdGenerator.generate()); + projections.emplace( + fieldSlots.back(), + sbe::makeE<sbe::EFunction>("getField"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(*_data.resultSlot), + sbe::makeE<sbe::EConstant>(std::string_view{ + field.c_str(), field.size()})))); + } + + return sbe::makeS<sbe::MakeObjStage>( + sbe::makeS<sbe::ProjectStage>(std::move(inputStage), std::move(projections)), + *_data.resultSlot, + boost::none, + std::vector<std::string>{}, + pn->proj.getRequiredFields(), + fieldSlots, + true, + false); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildProjectionDefault( + const QuerySolutionNode* root) { + using namespace std::literals; + + auto pn = static_cast<const ProjectionNodeDefault*>(root); + auto inputStage = build(pn->children[0]); + invariant(_data.resultSlot); + auto [slot, stage] = generateProjection( + &pn->proj, std::move(inputStage), &_slotIdGenerator, &_frameIdGenerator, *_data.resultSlot); + _data.resultSlot = slot; + return std::move(stage); +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildOr(const QuerySolutionNode* root) { + std::vector<std::unique_ptr<sbe::PlanStage>> inputStages; + std::vector<sbe::value::SlotVector> inputSlots; + + auto orn = static_cast<const OrNode*>(root); + + // Translate each child of the 'Or' node. Each child may produce new 'resultSlot' and + // recordIdSlot' stored in the _data member. We need to add these slots into the 'inputSlots' + // vector which is used as input to the union statge below. + for (auto&& child : orn->children) { + inputStages.push_back(build(child)); + invariant(_data.resultSlot); + invariant(_data.recordIdSlot); + inputSlots.push_back({*_data.resultSlot, *_data.recordIdSlot}); + } + + // Construct a union stage whose branches are translated children of the 'Or' node. + _data.resultSlot = _slotIdGenerator.generate(); + _data.recordIdSlot = _slotIdGenerator.generate(); + auto stage = sbe::makeS<sbe::UnionStage>(std::move(inputStages), + std::move(inputSlots), + sbe::makeSV(*_data.resultSlot, *_data.recordIdSlot)); + + if (orn->dedup) { + stage = sbe::makeS<sbe::HashAggStage>( + std::move(stage), sbe::makeSV(*_data.recordIdSlot), sbe::makeEM()); + } + + if (orn->filter) { + stage = generateFilter( + orn->filter.get(), std::move(stage), &_slotIdGenerator, *_data.resultSlot); + } + + return stage; +} + +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildText(const QuerySolutionNode* root) { + auto textNode = static_cast<const TextNode*>(root); + + invariant(_collection); + auto&& indexName = textNode->index.identifier.catalogName; + const auto desc = _collection->getIndexCatalog()->findIndexByName(_opCtx, indexName); + invariant(desc); + const auto accessMethod = static_cast<const FTSAccessMethod*>( + _collection->getIndexCatalog()->getEntry(desc)->accessMethod()); + invariant(accessMethod); + auto&& ftsSpec = accessMethod->getSpec(); + + // We assume here that node->ftsQuery is an FTSQueryImpl, not an FTSQueryNoop. In practice, this + // means that it is illegal to use the StageBuilder on a QuerySolution created by planning a + // query that contains "no-op" expressions. + auto ftsQuery = static_cast<fts::FTSQueryImpl&>(*textNode->ftsQuery); + + // A vector of the output slots for each index scan stage. Each stage outputs a record id and a + // record, so we expect each inner vector to be of length two. + std::vector<sbe::value::SlotVector> ixscanOutputSlots; + + const bool forward = true; + const bool inclusive = true; + auto makeKeyString = [&](const BSONObj& bsonKey) { + return std::make_unique<KeyString::Value>( + IndexEntryComparison::makeKeyStringFromBSONKeyForSeek( + bsonKey, + accessMethod->getSortedDataInterface()->getKeyStringVersion(), + accessMethod->getSortedDataInterface()->getOrdering(), + forward, + inclusive)); + }; + + std::vector<std::unique_ptr<sbe::PlanStage>> indexScanList; + for (const auto& term : ftsQuery.getTermsForBounds()) { + // TODO: Should we scan in the opposite direction? + auto startKeyBson = fts::FTSIndexFormat::getIndexKey( + 0, term, textNode->indexPrefix, ftsSpec.getTextIndexVersion()); + auto endKeyBson = fts::FTSIndexFormat::getIndexKey( + fts::MAX_WEIGHT, term, textNode->indexPrefix, ftsSpec.getTextIndexVersion()); + + auto recordSlot = _slotIdGenerator.generate(); + auto&& [recordIdSlot, ixscan] = + generateSingleIntervalIndexScan(_collection, + indexName, + forward, + makeKeyString(startKeyBson), + makeKeyString(endKeyBson), + recordSlot, + &_slotIdGenerator, + _yieldPolicy, + _data.trialRunProgressTracker.get()); + indexScanList.push_back(std::move(ixscan)); + ixscanOutputSlots.push_back(sbe::makeSV(recordIdSlot, recordSlot)); + } + + // Union will output a slot for the record id and another for the record. + _data.recordIdSlot = _slotIdGenerator.generate(); + auto unionRecordOutputSlot = _slotIdGenerator.generate(); + auto unionOutputSlots = sbe::makeSV(*_data.recordIdSlot, unionRecordOutputSlot); + + // Index scan output slots become the input slots to the union. + auto unionStage = + sbe::makeS<sbe::UnionStage>(std::move(indexScanList), ixscanOutputSlots, unionOutputSlots); + + // TODO: If text score metadata is requested, then we should sum over the text scores inside the + // index keys for a given document. This will require expression evaluation to be able to + // extract the score directly from the key string. + auto hashAggStage = sbe::makeS<sbe::HashAggStage>( + std::move(unionStage), sbe::makeSV(*_data.recordIdSlot), sbe::makeEM()); + + auto nljStage = makeLoopJoinForFetch(std::move(hashAggStage), *_data.recordIdSlot); + + // Add a special stage to apply 'ftsQuery' to matching documents, and then add a FilterStage to + // discard documents which do not match. + auto textMatchResultSlot = _slotIdGenerator.generate(); + auto textMatchStage = sbe::makeS<sbe::TextMatchStage>( + std::move(nljStage), ftsQuery, ftsSpec, *_data.resultSlot, textMatchResultSlot); + + // Filter based on the contents of the slot filled out by the TextMatchStage. + return sbe::makeS<sbe::FilterStage<false>>(std::move(textMatchStage), + sbe::makeE<sbe::EVariable>(textMatchResultSlot)); +} + +// Returns a non-null pointer to the root of a plan tree, or a non-OK status if the PlanStage tree +// could not be constructed. +std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::build(const QuerySolutionNode* root) { + static const stdx::unordered_map<StageType, + std::function<std::unique_ptr<sbe::PlanStage>( + SlotBasedStageBuilder&, const QuerySolutionNode* root)>> + kStageBuilders = { + {STAGE_COLLSCAN, std::mem_fn(&SlotBasedStageBuilder::buildCollScan)}, + {STAGE_IXSCAN, std::mem_fn(&SlotBasedStageBuilder::buildIndexScan)}, + {STAGE_FETCH, std::mem_fn(&SlotBasedStageBuilder::buildFetch)}, + {STAGE_LIMIT, std::mem_fn(&SlotBasedStageBuilder::buildLimit)}, + {STAGE_SKIP, std::mem_fn(&SlotBasedStageBuilder::buildSkip)}, + {STAGE_SORT_SIMPLE, std::mem_fn(&SlotBasedStageBuilder::buildSort)}, + {STAGE_SORT_DEFAULT, std::mem_fn(&SlotBasedStageBuilder::buildSort)}, + {STAGE_SORT_KEY_GENERATOR, std::mem_fn(&SlotBasedStageBuilder::buildSortKeyGeneraror)}, + {STAGE_PROJECTION_SIMPLE, std::mem_fn(&SlotBasedStageBuilder::buildProjectionSimple)}, + {STAGE_PROJECTION_DEFAULT, std::mem_fn(&SlotBasedStageBuilder::buildProjectionDefault)}, + {STAGE_OR, &SlotBasedStageBuilder::buildOr}, + {STAGE_TEXT, &SlotBasedStageBuilder::buildText}}; + + uassert(4822884, + str::stream() << "Can't build exec tree for node: " << root->toString(), + kStageBuilders.find(root->getType()) != kStageBuilders.end()); + + return std::invoke(kStageBuilders.at(root->getType()), *this, root); +} +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h new file mode 100644 index 00000000000..a91e1b8246a --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder.h @@ -0,0 +1,127 @@ +/** + * 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/trial_period_utils.h" +#include "mongo/db/exec/trial_run_progress_tracker.h" +#include "mongo/db/query/plan_yield_policy_sbe.h" +#include "mongo/db/query/stage_builder.h" + +namespace mongo::stage_builder { +/** + * Some auxiliary data returned by a 'SlotBasedStageBuilder' along with a PlanStage tree root, which + * is needed to execute the PlanStage tree. + */ +struct PlanStageData { + std::string debugString() const { + StringBuilder builder; + + if (resultSlot) { + builder << "$$RESULT=s" << *resultSlot << " "; + } + if (recordIdSlot) { + builder << "$$RID=s" << *recordIdSlot << " "; + } + if (oplogTsSlot) { + builder << "$$OPLOGTS=s" << *oplogTsSlot << " "; + } + + return builder.str(); + } + + boost::optional<sbe::value::SlotId> resultSlot; + boost::optional<sbe::value::SlotId> recordIdSlot; + boost::optional<sbe::value::SlotId> oplogTsSlot; + sbe::CompileCtx ctx; + bool shouldTrackLatestOplogTimestamp{false}; + bool shouldTrackResumeToken{false}; + // Used during the trial run of the runtime planner to track progress of the work done so far. + std::unique_ptr<TrialRunProgressTracker> trialRunProgressTracker; +}; + +/** + * A stage builder which builds an executable tree using slot-based PlanStages. + */ +class SlotBasedStageBuilder final : public StageBuilder<sbe::PlanStage> { +public: + SlotBasedStageBuilder(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + const QuerySolution& solution, + PlanYieldPolicySBE* yieldPolicy, + bool needsTrialRunProgressTracker) + : StageBuilder(opCtx, collection, cq, solution), _yieldPolicy(yieldPolicy) { + if (needsTrialRunProgressTracker) { + const auto maxNumResults{trial_period::getTrialPeriodNumToReturn(_cq)}; + const auto maxNumReads{trial_period::getTrialPeriodMaxWorks(_opCtx, _collection)}; + _data.trialRunProgressTracker = + std::make_unique<TrialRunProgressTracker>(maxNumResults, maxNumReads); + } + } + + std::unique_ptr<sbe::PlanStage> build(const QuerySolutionNode* root) final; + + PlanStageData getPlanStageData() { + return std::move(_data); + } + +private: + std::unique_ptr<sbe::PlanStage> buildCollScan(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildIndexScan(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildFetch(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildLimit(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildSkip(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildSort(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildSortKeyGeneraror(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildProjectionSimple(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildProjectionDefault(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildOr(const QuerySolutionNode* root); + std::unique_ptr<sbe::PlanStage> buildText(const QuerySolutionNode* root); + + std::unique_ptr<sbe::PlanStage> makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage, + sbe::value::SlotId recordIdKeySlot); + + sbe::value::SlotIdGenerator _slotIdGenerator; + sbe::value::FrameIdGenerator _frameIdGenerator; + sbe::value::SpoolIdGenerator _spoolIdGenerator; + + // If we have both limit and skip stages and the skip stage is beneath the limit, then we can + // combine these two stages into one. So, while processing the LIMIT stage we will save the + // limit value in this member and will handle it while processing the SKIP stage. + boost::optional<long long> _limit; + + PlanYieldPolicySBE* const _yieldPolicy; + + // Apart from generating just an execution tree, this builder will also produce some auxiliary + // data which is needed to execute the tree, such as a result slot, or a recordId slot. + PlanStageData _data; +}; +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp b/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp new file mode 100644 index 00000000000..fbfcef2cca6 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp @@ -0,0 +1,362 @@ +/** + * 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/query/sbe_stage_builder_coll_scan.h" + +#include "mongo/db/catalog/collection.h" +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/exchange.h" +#include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/stages/scan.h" +#include "mongo/db/query/sbe_stage_builder_filter.h" +#include "mongo/db/query/util/make_data_structure.h" +#include "mongo/db/storage/oplog_hack.h" +#include "mongo/logv2/log.h" +#include "mongo/util/str.h" + +namespace mongo::stage_builder { +namespace { +/** + * Checks whether a callback function should be created for a ScanStage and returns it, if so. The + * logic in the provided callback will be executed when the ScanStage is opened or reopened. + */ +sbe::ScanOpenCallback makeOpenCallbackIfNeeded(const Collection* collection, + const CollectionScanNode* csn) { + if (csn->direction == CollectionScanParams::FORWARD && csn->shouldWaitForOplogVisibility) { + invariant(!csn->tailable); + invariant(collection->ns().isOplog()); + + return [](OperationContext* opCtx, const Collection* collection, bool reOpen) { + if (!reOpen) { + // Forward, non-tailable scans from the oplog need to wait until all oplog entries + // before the read begins to be visible. This isn't needed for reverse scans because + // we only hide oplog entries from forward scans, and it isn't necessary for tailing + // cursors because they ignore EOF and will eventually see all writes. Forward, + // non-tailable scans are the only case where a meaningful EOF will be seen that + // might not include writes that finished before the read started. This also must be + // done before we create the cursor as that is when we establish the endpoint for + // the cursor. Also call abandonSnapshot to make sure that we are using a fresh + // storage engine snapshot while waiting. Otherwise, we will end up reading from the + // snapshot where the oplog entries are not yet visible even after the wait. + + opCtx->recoveryUnit()->abandonSnapshot(); + collection->getRecordStore()->waitForAllEarlierOplogWritesToBeVisible(opCtx); + } + }; + } + return {}; +} + +/** + * Creates a collection scan sub-tree optimized for oplog scans. We can built an optimized scan + * when there is a predicted on the 'ts' field of the oplog collection. + * + * 1. If a lower bound on 'ts' is present, the collection scan will seek directly to the RecordId + * of an oplog entry as close to this lower bound as possible without going higher. + * 1.1 If the query is just a lower bound on 'ts' on a forward scan, every document in the + * collection after the first matching one must also match. To avoid wasting time + * running the filter on every document to be returned, we will stop applying the filter + * once it finds the first match. + * 2. If an upper bound on 'ts' is present, the collection scan will stop and return EOF the first + * time it fetches a document that does not pass the filter and has 'ts' greater than the upper + * bound. + */ +std::tuple<sbe::value::SlotId, + sbe::value::SlotId, + boost::optional<sbe::value::SlotId>, + std::unique_ptr<sbe::PlanStage>> +generateOptimizedOplogScan(OperationContext* opCtx, + const Collection* collection, + const CollectionScanNode* csn, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + invariant(collection->ns().isOplog()); + // The minTs and maxTs optimizations are not compatible with resumeAfterRecordId and can only + // be done for a forward scan. + invariant(!csn->resumeAfterRecordId); + invariant(csn->direction == CollectionScanParams::FORWARD); + + auto resultSlot = slotIdGenerator->generate(); + auto recordIdSlot = slotIdGenerator->generate(); + + // See if the RecordStore supports the oplogStartHack. If so, the scan will start from the + // RecordId stored in seekRecordId. + auto [seekRecordId, seekRecordIdSlot] = + [&]() -> std::pair<boost::optional<RecordId>, boost::optional<sbe::value::SlotId>> { + if (csn->minTs) { + auto goal = oploghack::keyForOptime(*csn->minTs); + if (goal.isOK()) { + auto startLoc = + collection->getRecordStore()->oplogStartHack(opCtx, goal.getValue()); + if (startLoc && !startLoc->isNull()) { + LOGV2_DEBUG(205841, 3, "Using direct oplog seek"); + return {startLoc, slotIdGenerator->generate()}; + } + } + } + return {}; + }(); + + // Check if we need to project out an oplog 'ts' field as part of the collection scan. We will + // need it either when 'maxTs' bound has been provided, so that we can apply an EOF filter, of + // if we need to track the latest oplog timestamp. + auto [fields, slots, tsSlot] = [&]() -> std::tuple<std::vector<std::string>, + sbe::value::SlotVector, + boost::optional<sbe::value::SlotId>> { + // Don't project the 'ts' if stopApplyingFilterAfterFirstMatch is 'true'. We will have + // another scan stage where it will be done. + if (!csn->stopApplyingFilterAfterFirstMatch && + (csn->maxTs || csn->shouldTrackLatestOplogTimestamp)) { + auto tsSlot = slotIdGenerator->generate(); + return {{repl::OpTime::kTimestampFieldName}, sbe::makeSV(tsSlot), tsSlot}; + } + return {}; + }(); + + NamespaceStringOrUUID nss{collection->ns().db().toString(), collection->uuid()}; + auto stage = sbe::makeS<sbe::ScanStage>(nss, + resultSlot, + recordIdSlot, + std::move(fields), + std::move(slots), + seekRecordIdSlot, + true /* forward */, + yieldPolicy, + tracker, + makeOpenCallbackIfNeeded(collection, csn)); + + // Start the scan from the seekRecordId if we can use the oplogStartHack. + if (seekRecordId) { + invariant(seekRecordIdSlot); + + // Project the start RecordId as a seekRecordIdSlot and feed it to the inner side (scan). + stage = sbe::makeS<sbe::LoopJoinStage>( + sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + *seekRecordIdSlot, + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt64, + seekRecordId->repr())), + std::move(stage), + sbe::makeSV(), + sbe::makeSV(*seekRecordIdSlot), + nullptr); + } + + // Add an EOF filter to stop the scan after we fetch the first document that has 'ts' greater + // than the upper bound. + if (csn->maxTs) { + // The 'maxTs' optimization is not compatible with 'stopApplyingFilterAfterFirstMatch'. + invariant(!csn->stopApplyingFilterAfterFirstMatch); + invariant(tsSlot); + + stage = sbe::makeS<sbe::FilterStage<false, true>>( + std::move(stage), + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::lessEq, + sbe::makeE<sbe::EVariable>(*tsSlot), + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Timestamp, + (*csn->maxTs).asULL()))); + } + + if (csn->filter) { + stage = generateFilter(csn->filter.get(), std::move(stage), slotIdGenerator, resultSlot); + + // We may be requested to stop applying the filter after the first match. This can happen + // if the query is just a lower bound on 'ts' on a forward scan. In this case every document + // in the collection after the first matching one must also match, so there is no need to + // run the filter on such elements. + // + // To apply this optimization we will construct the following sub-tree: + // + // nlj [] [seekRecordIdSlot] + // left + // limit 1 + // filter <predicate> + // <stage> + // right + // seek seekRecordIdSlot resultSlot recordIdSlot @coll + // + // Here, the nested loop join outer branch is the collection scan we constructed above, with + // a csn->filter predicate sitting on top. The 'limit 1' stage is to ensure this branch + // returns a single row. Once executed, this branch will filter out documents which doesn't + // satisfy the predicate, and will return the first document, along with a RecordId, that + // matches. This RecordId is then used as a starting point of the collection scan in the + // inner branch, and the execution will continue from this point further on, without + // applying the filter. + if (csn->stopApplyingFilterAfterFirstMatch) { + std::tie(fields, slots, tsSlot) = + [&]() -> std::tuple<std::vector<std::string>, + sbe::value::SlotVector, + boost::optional<sbe::value::SlotId>> { + if (csn->shouldTrackLatestOplogTimestamp) { + auto tsSlot = slotIdGenerator->generate(); + return {{repl::OpTime::kTimestampFieldName}, sbe::makeSV(tsSlot), tsSlot}; + } + return {}; + }(); + + seekRecordIdSlot = recordIdSlot; + resultSlot = slotIdGenerator->generate(); + recordIdSlot = slotIdGenerator->generate(); + + stage = sbe::makeS<sbe::LoopJoinStage>( + sbe::makeS<sbe::LimitSkipStage>(std::move(stage), 1, boost::none), + sbe::makeS<sbe::ScanStage>(nss, + resultSlot, + recordIdSlot, + std::move(fields), + std::move(slots), + seekRecordIdSlot, + true /* forward */, + yieldPolicy, + tracker), + sbe::makeSV(), + sbe::makeSV(*seekRecordIdSlot), + nullptr); + } + } + + return {resultSlot, + recordIdSlot, + csn->shouldTrackLatestOplogTimestamp ? tsSlot : boost::none, + std::move(stage)}; +} + +/** + * Generates a generic collecion scan sub-tree. If a resume token has been provided, the scan will + * start from a RecordId contained within this token, otherwise from the beginning of the + * collection. + */ +std::tuple<sbe::value::SlotId, + sbe::value::SlotId, + boost::optional<sbe::value::SlotId>, + std::unique_ptr<sbe::PlanStage>> +generateGenericCollScan(const Collection* collection, + const CollectionScanNode* csn, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + const auto forward = csn->direction == CollectionScanParams::FORWARD; + + auto resultSlot = slotIdGenerator->generate(); + auto recordIdSlot = slotIdGenerator->generate(); + auto seekRecordIdSlot = boost::make_optional(static_cast<bool>(csn->resumeAfterRecordId), + slotIdGenerator->generate()); + + // See if we need to project out an oplog latest timestamp. + auto [fields, slots, tsSlot] = [&]() -> std::tuple<std::vector<std::string>, + sbe::value::SlotVector, + boost::optional<sbe::value::SlotId>> { + if (csn->shouldTrackLatestOplogTimestamp) { + invariant(collection->ns().isOplog()); + + auto tsSlot = slotIdGenerator->generate(); + return {{repl::OpTime::kTimestampFieldName}, sbe::makeSV(tsSlot), tsSlot}; + } + return {}; + }(); + + NamespaceStringOrUUID nss{collection->ns().db().toString(), collection->uuid()}; + auto stage = sbe::makeS<sbe::ScanStage>(nss, + resultSlot, + recordIdSlot, + std::move(fields), + std::move(slots), + seekRecordIdSlot, + forward, + yieldPolicy, + tracker, + makeOpenCallbackIfNeeded(collection, csn)); + + // Check if the scan should be started after the provided resume RecordId and construct a nested + // loop join sub-tree to project out the resume RecordId as a seekRecordIdSlot and feed it to + // the inner side (scan). + // + // Note that we also inject a 'skip 1' stage on top of the inner branch, as we need to start + // _after_ the resume RecordId. + // + // TODO SERVER-48472: raise KeyNotFound error if we cannot position the cursor on + // seekRecordIdSlot. + if (seekRecordIdSlot) { + stage = sbe::makeS<sbe::LoopJoinStage>( + sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + *seekRecordIdSlot, + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt64, + csn->resumeAfterRecordId->repr())), + sbe::makeS<sbe::LimitSkipStage>(std::move(stage), boost::none, 1), + sbe::makeSV(), + sbe::makeSV(*seekRecordIdSlot), + nullptr); + } + + if (csn->filter) { + // The 'stopApplyingFilterAfterFirstMatch' optimization is only applicable when the 'ts' + // lower bound is also provided for an oplog scan, and is handled in + // 'generateOptimizedOplogScan()'. + invariant(!csn->stopApplyingFilterAfterFirstMatch); + + stage = generateFilter(csn->filter.get(), std::move(stage), slotIdGenerator, resultSlot); + } + + return {resultSlot, recordIdSlot, tsSlot, std::move(stage)}; +} +} // namespace + +std::tuple<sbe::value::SlotId, + sbe::value::SlotId, + boost::optional<sbe::value::SlotId>, + std::unique_ptr<sbe::PlanStage>> +generateCollScan(OperationContext* opCtx, + const Collection* collection, + const CollectionScanNode* csn, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + uassert(4822889, "Tailable collection scans are not supported in SBE", !csn->tailable); + + auto [resultSlot, recordIdSlot, oplogTsSlot, stage] = [&]() { + if (csn->minTs || csn->maxTs) { + return generateOptimizedOplogScan( + opCtx, collection, csn, slotIdGenerator, yieldPolicy, tracker); + } else { + return generateGenericCollScan(collection, csn, slotIdGenerator, yieldPolicy, tracker); + } + }(); + + return {resultSlot, recordIdSlot, oplogTsSlot, std::move(stage)}; +} +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_coll_scan.h b/src/mongo/db/query/sbe_stage_builder_coll_scan.h new file mode 100644 index 00000000000..f9d91f5e550 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_coll_scan.h @@ -0,0 +1,59 @@ +/** + * 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/exec/trial_run_progress_tracker.h" +#include "mongo/db/query/query_solution.h" + +namespace mongo::stage_builder { +/** + * Generates an SBE plan stage sub-tree implementing an collection scan. + * + * On success, a tuple containing the following data is returned: + * * A slot to access a fetched document (a resultSlot) + * * A slot to access a recordId (a recordIdSlot) + * * An optional slot to access a latest oplog timestamp (oplogTsSlot), if we scan the oplog and + * were requested to track this data. + * * A generated PlanStage sub-tree. + * + * In cases of an error, throws. + */ +std::tuple<sbe::value::SlotId, + sbe::value::SlotId, + boost::optional<sbe::value::SlotId>, + std::unique_ptr<sbe::PlanStage>> +generateCollScan(OperationContext* opCtx, + const Collection* collection, + const CollectionScanNode* csn, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker); +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp new file mode 100644 index 00000000000..cb349949adc --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp @@ -0,0 +1,1341 @@ +/** + * 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/query/sbe_stage_builder_expression.h" + +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/stages/traverse.h" +#include "mongo/db/exec/sbe/stages/union.h" +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/pipeline/accumulator.h" +#include "mongo/db/pipeline/expression_visitor.h" +#include "mongo/db/pipeline/expression_walker.h" +#include "mongo/db/query/projection_parser.h" +#include "mongo/util/str.h" + +namespace mongo::stage_builder { +namespace { +std::pair<sbe::value::TypeTags, sbe::value::Value> convertFrom(Value val) { + // TODO: Either make this conversion unnecessary by changing the value representation in + // ExpressionConstant, or provide a nicer way to convert directly from Document/Value to + // sbe::Value. + BSONObjBuilder bob; + val.addToBsonObj(&bob, ""_sd); + auto obj = bob.done(); + auto be = obj.objdata(); + auto end = be + sbe::value::readFromMemory<uint32_t>(be); + return sbe::bson::convertFrom(false, be + 4, end, 0); +} + +struct ExpressionVisitorContext { + struct VarsFrame { + std::deque<Variables::Id> variablesToBind; + sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> boundVariables; + + template <class... Args> + VarsFrame(Args&&... args) + : variablesToBind{std::forward<Args>(args)...}, boundVariables{} {} + }; + + struct LogicalExpressionEvalFrame { + std::unique_ptr<sbe::PlanStage> savedTraverseStage; + sbe::value::SlotVector savedRelevantSlots; + + sbe::value::SlotId nextBranchResultSlot; + + std::vector<std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>>> branches; + + LogicalExpressionEvalFrame(std::unique_ptr<sbe::PlanStage> traverseStage, + const sbe::value::SlotVector& relevantSlots, + sbe::value::SlotId nextBranchResultSlot) + : savedTraverseStage(std::move(traverseStage)), + savedRelevantSlots(relevantSlots), + nextBranchResultSlot(nextBranchResultSlot) {} + }; + + ExpressionVisitorContext(std::unique_ptr<sbe::PlanStage> inputStage, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::FrameIdGenerator* frameIdGenerator, + sbe::value::SlotId rootSlot, + sbe::value::SlotVector* relevantSlots) + : traverseStage(std::move(inputStage)), + slotIdGenerator(slotIdGenerator), + frameIdGenerator(frameIdGenerator), + rootSlot(rootSlot), + relevantSlots(relevantSlots) {} + + void ensureArity(size_t arity) { + invariant(exprs.size() >= arity); + } + + std::unique_ptr<sbe::EExpression> popExpr() { + auto expr = std::move(exprs.top()); + exprs.pop(); + return expr; + } + + void pushExpr(std::unique_ptr<sbe::EExpression> expr) { + exprs.push(std::move(expr)); + } + + void pushExpr(std::unique_ptr<sbe::EExpression> expr, std::unique_ptr<sbe::PlanStage> stage) { + exprs.push(std::move(expr)); + traverseStage = std::move(stage); + } + + /** + * Temporarily reset 'traverseStage' and 'relevantSlots' so they are prepared for translating a + * $and/$or branch. (They will be restored later using the saved values in the + * 'logicalExpressionEvalFrameStack' top entry.) The new 'traverseStage' is actually a + * projection that will evaluate to a constant false (for $and) or true (for $or). Once this + * branch is fully contructed, it will have a filter stage that will either filter out the + * constant (when branch evaluation does not short circuit) or produce the constant value + * (therefore producing the short-circuit result). These branches are part of a union stage, so + * each time a branch fails to produce a value, execution moves on to the next branch. A limit + * stage above the union ensures that execution halts once one of the branches produces a + * result. + */ + void prepareToTranslateShortCircuitingBranch(sbe::EPrimBinary::Op logicOp, + sbe::value::SlotId branchResultSlot) { + invariant(!logicalExpressionEvalFrameStack.empty()); + logicalExpressionEvalFrameStack.top().nextBranchResultSlot = branchResultSlot; + + auto shortCircuitVal = (logicOp == sbe::EPrimBinary::logicOr); + auto shortCircuitExpr = sbe::makeE<sbe::EConstant>( + sbe::value::TypeTags::Boolean, sbe::value::bitcastFrom(shortCircuitVal)); + traverseStage = sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + branchResultSlot, + std::move(shortCircuitExpr)); + + // Slots created in a previous branch for this $and/$or are not accessible to any stages in + // this new branch, so we clear them from the 'relevantSlots' list. + *relevantSlots = logicalExpressionEvalFrameStack.top().savedRelevantSlots; + + // The 'branchResultSlot' is where the new branch will place its result in the event of a + // short circuit, and it must be visible to the union stage after the branch executes. + relevantSlots->push_back(branchResultSlot); + } + + /** + * This does the same thing as 'prepareToTranslateShortCircuitingBranch' but is intended to the + * last branch in an $and/$or, which cannot short circuit. + */ + void prepareToTranslateConcludingLogicalBranch() { + invariant(!logicalExpressionEvalFrameStack.empty()); + + traverseStage = sbe::makeS<sbe::CoScanStage>(); + *relevantSlots = logicalExpressionEvalFrameStack.top().savedRelevantSlots; + } + + std::tuple<sbe::value::SlotId, + std::unique_ptr<sbe::EExpression>, + std::unique_ptr<sbe::PlanStage>> + done() { + invariant(exprs.size() == 1); + return {slotIdGenerator->generate(), popExpr(), std::move(traverseStage)}; + } + + std::unique_ptr<sbe::PlanStage> traverseStage; + sbe::value::SlotIdGenerator* slotIdGenerator; + sbe::value::FrameIdGenerator* frameIdGenerator; + sbe::value::SlotId rootSlot; + std::stack<std::unique_ptr<sbe::EExpression>> exprs; + std::map<Variables::Id, sbe::value::SlotId> environment; + std::stack<VarsFrame> varsFrameStack; + std::stack<LogicalExpressionEvalFrame> logicalExpressionEvalFrameStack; + // See the comment above the generateExpression() declaration for an explanation of the + // 'relevantSlots' list. + sbe::value::SlotVector* relevantSlots; +}; + +/** + * Generate an EExpression that converts a value (contained in a variable bound to 'branchRef') that + * can be of any type to a Boolean value based on MQL's definition of truth for the branch of a + * "$and" or "$or" expression. + */ +std::unique_ptr<sbe::EExpression> generateExpressionForLogicBranch(sbe::EVariable branchRef) { + // Make an expression that compares the value in 'branchRef' to the result of evaluating the + // 'valExpr' expression. The comparison uses cmp3w, so that can handle comparisons between + // values with different types. + auto makeNeqCheck = [&branchRef](std::unique_ptr<sbe::EExpression> valExpr) { + return sbe::makeE<sbe::EPrimBinary>( + sbe::EPrimBinary::neq, + sbe::makeE<sbe::EPrimBinary>( + sbe::EPrimBinary::cmp3w, branchRef.clone(), std::move(valExpr)), + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt64, 0)); + }; + + // If any of these are false, the branch is considered false for the purposes of the + // $and/$or. + auto checkExists = sbe::makeE<sbe::EFunction>("exists", sbe::makeEs(branchRef.clone())); + auto checkNotNull = sbe::makeE<sbe::EPrimUnary>( + sbe::EPrimUnary::logicNot, + sbe::makeE<sbe::EFunction>("isNull", sbe::makeEs(branchRef.clone()))); + auto checkNotFalse = + makeNeqCheck(sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Boolean, false)); + auto checkNotZero = + makeNeqCheck(sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt64, 0)); + + return sbe::makeE<sbe::EPrimBinary>( + sbe::EPrimBinary::logicAnd, + std::move(checkExists), + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicAnd, + std::move(checkNotNull), + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicAnd, + std::move(checkNotFalse), + std::move(checkNotZero)))); +} + +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateTraverseHelper( + std::unique_ptr<sbe::PlanStage> inputStage, + sbe::value::SlotId inputSlot, + const FieldPath& fp, + size_t level, + sbe::value::SlotIdGenerator* slotIdGenerator) { + using namespace std::literals; + + invariant(level < fp.getPathLength()); + + // The field we will be traversing at the current nested level. + auto fieldSlot{slotIdGenerator->generate()}; + // The result coming from the 'in' branch of the traverse plan stage. + auto outputSlot{slotIdGenerator->generate()}; + + // Generate the projection stage to read a sub-field at the current nested level and bind it + // to 'fieldSlot'. + inputStage = sbe::makeProjectStage( + std::move(inputStage), + fieldSlot, + sbe::makeE<sbe::EFunction>( + "getField"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(inputSlot), sbe::makeE<sbe::EConstant>([&]() { + auto fieldName = fp.getFieldName(level); + return std::string_view{fieldName.rawData(), fieldName.size()}; + }())))); + + std::unique_ptr<sbe::PlanStage> innerBranch; + if (level == fp.getPathLength() - 1) { + innerBranch = sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + outputSlot, + sbe::makeE<sbe::EVariable>(fieldSlot)); + } else { + // Generate nested traversal. + auto [slot, stage] = generateTraverseHelper( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + fieldSlot, + fp, + level + 1, + slotIdGenerator); + innerBranch = + sbe::makeProjectStage(std::move(stage), outputSlot, sbe::makeE<sbe::EVariable>(slot)); + } + + // The final traverse stage for the current nested level. + return {outputSlot, + sbe::makeS<sbe::TraverseStage>(std::move(inputStage), + std::move(innerBranch), + fieldSlot, + outputSlot, + outputSlot, + sbe::makeSV(), + nullptr, + nullptr, + 1)}; +} + +/** + * For the given MatchExpression 'expr', generates a path traversal SBE plan stage sub-tree + * implementing the comparison expression. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateTraverse( + std::unique_ptr<sbe::PlanStage> inputStage, + sbe::value::SlotId inputSlot, + bool expectsDocumentInputOnly, + const FieldPath& fp, + sbe::value::SlotIdGenerator* slotIdGenerator) { + if (expectsDocumentInputOnly) { + // When we know for sure that 'inputSlot' will be a document and _not_ an array (such as + // when traversing the root document), we can generate a simpler expression. + return generateTraverseHelper(std::move(inputStage), inputSlot, fp, 0, slotIdGenerator); + } else { + // The general case: the value in the 'inputSlot' may be an array that will require + // traversal. + auto outputSlot{slotIdGenerator->generate()}; + auto [innerBranchOutputSlot, innerBranch] = generateTraverseHelper( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + inputSlot, + fp, + 0, // level + slotIdGenerator); + return {outputSlot, + sbe::makeS<sbe::TraverseStage>(std::move(inputStage), + std::move(innerBranch), + inputSlot, + outputSlot, + innerBranchOutputSlot, + sbe::makeSV(), + nullptr, + nullptr)}; + } +} + +class ExpressionPreVisitor final : public ExpressionVisitor { +public: + ExpressionPreVisitor(ExpressionVisitorContext* context) : _context{context} {} + + void visit(ExpressionConstant* expr) final {} + void visit(ExpressionAbs* expr) final {} + void visit(ExpressionAdd* expr) final {} + void visit(ExpressionAllElementsTrue* expr) final {} + void visit(ExpressionAnd* expr) final { + visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicAnd); + } + void visit(ExpressionAnyElementTrue* expr) final {} + void visit(ExpressionArray* expr) final {} + void visit(ExpressionArrayElemAt* expr) final {} + void visit(ExpressionFirst* expr) final {} + void visit(ExpressionLast* expr) final {} + void visit(ExpressionObjectToArray* expr) final {} + void visit(ExpressionArrayToObject* expr) final {} + void visit(ExpressionBsonSize* expr) final {} + void visit(ExpressionCeil* expr) final {} + void visit(ExpressionCoerceToBool* expr) final {} + void visit(ExpressionCompare* expr) final {} + void visit(ExpressionConcat* expr) final {} + void visit(ExpressionConcatArrays* expr) final {} + void visit(ExpressionCond* expr) final {} + void visit(ExpressionDateFromString* expr) final {} + void visit(ExpressionDateFromParts* expr) final {} + void visit(ExpressionDateToParts* expr) final {} + void visit(ExpressionDateToString* expr) final {} + void visit(ExpressionDivide* expr) final {} + void visit(ExpressionExp* expr) final {} + void visit(ExpressionFieldPath* expr) final {} + void visit(ExpressionFilter* expr) final {} + void visit(ExpressionFloor* expr) final {} + void visit(ExpressionIfNull* expr) final {} + void visit(ExpressionIn* expr) final {} + void visit(ExpressionIndexOfArray* expr) final {} + void visit(ExpressionIndexOfBytes* expr) final {} + void visit(ExpressionIndexOfCP* expr) final {} + void visit(ExpressionIsNumber* expr) final {} + void visit(ExpressionLet* expr) final { + _context->varsFrameStack.push(ExpressionVisitorContext::VarsFrame{ + std::begin(expr->getOrderedVariableIds()), std::end(expr->getOrderedVariableIds())}); + } + void visit(ExpressionLn* expr) final {} + void visit(ExpressionLog* expr) final {} + void visit(ExpressionLog10* expr) final {} + void visit(ExpressionMap* expr) final {} + void visit(ExpressionMeta* expr) final {} + void visit(ExpressionMod* expr) final {} + void visit(ExpressionMultiply* expr) final {} + void visit(ExpressionNot* expr) final {} + void visit(ExpressionObject* expr) final {} + void visit(ExpressionOr* expr) final { + visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicOr); + } + void visit(ExpressionPow* expr) final {} + void visit(ExpressionRange* expr) final {} + void visit(ExpressionReduce* expr) final {} + void visit(ExpressionReplaceOne* expr) final {} + void visit(ExpressionReplaceAll* expr) final {} + void visit(ExpressionSetDifference* expr) final {} + void visit(ExpressionSetEquals* expr) final {} + void visit(ExpressionSetIntersection* expr) final {} + void visit(ExpressionSetIsSubset* expr) final {} + void visit(ExpressionSetUnion* expr) final {} + void visit(ExpressionSize* expr) final {} + void visit(ExpressionReverseArray* expr) final {} + void visit(ExpressionSlice* expr) final {} + void visit(ExpressionIsArray* expr) final {} + void visit(ExpressionRound* expr) final {} + void visit(ExpressionSplit* expr) final {} + void visit(ExpressionSqrt* expr) final {} + void visit(ExpressionStrcasecmp* expr) final {} + void visit(ExpressionSubstrBytes* expr) final {} + void visit(ExpressionSubstrCP* expr) final {} + void visit(ExpressionStrLenBytes* expr) final {} + void visit(ExpressionBinarySize* expr) final {} + void visit(ExpressionStrLenCP* expr) final {} + void visit(ExpressionSubtract* expr) final {} + void visit(ExpressionSwitch* expr) final {} + void visit(ExpressionToLower* expr) final {} + void visit(ExpressionToUpper* expr) final {} + void visit(ExpressionTrim* expr) final {} + void visit(ExpressionTrunc* expr) final {} + void visit(ExpressionType* expr) final {} + void visit(ExpressionZip* expr) final {} + void visit(ExpressionConvert* expr) final {} + void visit(ExpressionRegexFind* expr) final {} + void visit(ExpressionRegexFindAll* expr) final {} + void visit(ExpressionRegexMatch* expr) final {} + void visit(ExpressionCosine* expr) final {} + void visit(ExpressionSine* expr) final {} + void visit(ExpressionTangent* expr) final {} + void visit(ExpressionArcCosine* expr) final {} + void visit(ExpressionArcSine* expr) final {} + void visit(ExpressionArcTangent* expr) final {} + void visit(ExpressionArcTangent2* expr) final {} + void visit(ExpressionHyperbolicArcTangent* expr) final {} + void visit(ExpressionHyperbolicArcCosine* expr) final {} + void visit(ExpressionHyperbolicArcSine* expr) final {} + void visit(ExpressionHyperbolicTangent* expr) final {} + void visit(ExpressionHyperbolicCosine* expr) final {} + void visit(ExpressionHyperbolicSine* expr) final {} + void visit(ExpressionDegreesToRadians* expr) final {} + void visit(ExpressionRadiansToDegrees* expr) final {} + void visit(ExpressionDayOfMonth* expr) final {} + void visit(ExpressionDayOfWeek* expr) final {} + void visit(ExpressionDayOfYear* expr) final {} + void visit(ExpressionHour* expr) final {} + void visit(ExpressionMillisecond* expr) final {} + void visit(ExpressionMinute* expr) final {} + void visit(ExpressionMonth* expr) final {} + void visit(ExpressionSecond* expr) final {} + void visit(ExpressionWeek* expr) final {} + void visit(ExpressionIsoWeekYear* expr) final {} + void visit(ExpressionIsoDayOfWeek* expr) final {} + void visit(ExpressionIsoWeek* expr) final {} + void visit(ExpressionYear* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorAvg>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorMax>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorMin>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorStdDevPop>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorStdDevSamp>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorSum>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorMergeObjects>* expr) final {} + void visit(ExpressionTests::Testable* expr) final {} + void visit(ExpressionInternalJsEmit* expr) final {} + void visit(ExpressionInternalFindSlice* expr) final {} + void visit(ExpressionInternalFindPositional* expr) final {} + void visit(ExpressionInternalFindElemMatch* expr) final {} + void visit(ExpressionFunction* expr) final {} + void visit(ExpressionInternalRemoveFieldTombstones* expr) final {} + void visit(ExpressionRandom* expr) final {} + +private: + void visitMultiBranchLogicExpression(Expression* expr, sbe::EPrimBinary::Op logicOp) { + invariant(logicOp == sbe::EPrimBinary::logicOr || logicOp == sbe::EPrimBinary::logicAnd); + + if (expr->getChildren().size() < 2) { + // All this bookkeeping is only necessary for short circuiting, so we can skip it if we + // don't have two or more branches. + return; + } + + auto branchResultSlot = _context->slotIdGenerator->generate(); + _context->logicalExpressionEvalFrameStack.emplace( + std::move(_context->traverseStage), *_context->relevantSlots, branchResultSlot); + + _context->prepareToTranslateShortCircuitingBranch(logicOp, branchResultSlot); + } + + ExpressionVisitorContext* _context; +}; + +class ExpressionInVisitor final : public ExpressionVisitor { +public: + ExpressionInVisitor(ExpressionVisitorContext* context) : _context{context} {} + + void visit(ExpressionConstant* expr) final {} + void visit(ExpressionAbs* expr) final {} + void visit(ExpressionAdd* expr) final {} + void visit(ExpressionAllElementsTrue* expr) final {} + void visit(ExpressionAnd* expr) final { + visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicAnd); + } + void visit(ExpressionAnyElementTrue* expr) final {} + void visit(ExpressionArray* expr) final {} + void visit(ExpressionArrayElemAt* expr) final {} + void visit(ExpressionFirst* expr) final {} + void visit(ExpressionLast* expr) final {} + void visit(ExpressionObjectToArray* expr) final {} + void visit(ExpressionArrayToObject* expr) final {} + void visit(ExpressionBsonSize* expr) final {} + void visit(ExpressionCeil* expr) final {} + void visit(ExpressionCoerceToBool* expr) final {} + void visit(ExpressionCompare* expr) final {} + void visit(ExpressionConcat* expr) final {} + void visit(ExpressionConcatArrays* expr) final {} + void visit(ExpressionCond* expr) final {} + void visit(ExpressionDateFromString* expr) final {} + void visit(ExpressionDateFromParts* expr) final {} + void visit(ExpressionDateToParts* expr) final {} + void visit(ExpressionDateToString* expr) final {} + void visit(ExpressionDivide* expr) final {} + void visit(ExpressionExp* expr) final {} + void visit(ExpressionFieldPath* expr) final {} + void visit(ExpressionFilter* expr) final {} + void visit(ExpressionFloor* expr) final {} + void visit(ExpressionIfNull* expr) final {} + void visit(ExpressionIn* expr) final {} + void visit(ExpressionIndexOfArray* expr) final {} + void visit(ExpressionIndexOfBytes* expr) final {} + void visit(ExpressionIndexOfCP* expr) final {} + void visit(ExpressionIsNumber* expr) final {} + void visit(ExpressionLet* expr) final { + // This visitor fires after each variable definition in a $let expression. The top of the + // _context's expression stack will be an expression defining the variable initializer. We + // use a separate frame stack ('varsFrameStack') to keep track of which variable we are + // visiting, so we can appropriately bind the initializer. + invariant(!_context->varsFrameStack.empty()); + auto& currentFrame = _context->varsFrameStack.top(); + + invariant(!currentFrame.variablesToBind.empty()); + + auto varToBind = currentFrame.variablesToBind.front(); + currentFrame.variablesToBind.pop_front(); + + // We create two bindings. First, the initializer result is bound to a slot when this + // ProjectStage executes. + auto slotToBind = _context->slotIdGenerator->generate(); + invariant(currentFrame.boundVariables.find(slotToBind) == + currentFrame.boundVariables.end()); + _context->ensureArity(1); + currentFrame.boundVariables.insert(std::make_pair(slotToBind, _context->popExpr())); + + // Second, we bind this variables AST-level name (with type Variable::Id) to the SlotId that + // will be used for compilation and execution. Once this "stage builder" finishes, these + // Variable::Id bindings will no longer be relevant. + invariant(_context->environment.find(varToBind) == _context->environment.end()); + _context->environment.insert({varToBind, slotToBind}); + + if (currentFrame.variablesToBind.empty()) { + // We have traversed every variable initializer, and this is the last "infix" visit for + // this $let. Add the the ProjectStage that will perform the actual binding now, so that + // it executes before the "in" portion of the $let statement does. + _context->traverseStage = sbe::makeS<sbe::ProjectStage>( + std::move(_context->traverseStage), std::move(currentFrame.boundVariables)); + } + } + void visit(ExpressionLn* expr) final {} + void visit(ExpressionLog* expr) final {} + void visit(ExpressionLog10* expr) final {} + void visit(ExpressionMap* expr) final {} + void visit(ExpressionMeta* expr) final {} + void visit(ExpressionMod* expr) final {} + void visit(ExpressionMultiply* expr) final {} + void visit(ExpressionNot* expr) final {} + void visit(ExpressionObject* expr) final {} + void visit(ExpressionOr* expr) final { + visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicOr); + } + void visit(ExpressionPow* expr) final {} + void visit(ExpressionRange* expr) final {} + void visit(ExpressionReduce* expr) final {} + void visit(ExpressionReplaceOne* expr) final {} + void visit(ExpressionReplaceAll* expr) final {} + void visit(ExpressionSetDifference* expr) final {} + void visit(ExpressionSetEquals* expr) final {} + void visit(ExpressionSetIntersection* expr) final {} + void visit(ExpressionSetIsSubset* expr) final {} + void visit(ExpressionSetUnion* expr) final {} + void visit(ExpressionSize* expr) final {} + void visit(ExpressionReverseArray* expr) final {} + void visit(ExpressionSlice* expr) final {} + void visit(ExpressionIsArray* expr) final {} + void visit(ExpressionRound* expr) final {} + void visit(ExpressionSplit* expr) final {} + void visit(ExpressionSqrt* expr) final {} + void visit(ExpressionStrcasecmp* expr) final {} + void visit(ExpressionSubstrBytes* expr) final {} + void visit(ExpressionSubstrCP* expr) final {} + void visit(ExpressionStrLenBytes* expr) final {} + void visit(ExpressionBinarySize* expr) final {} + void visit(ExpressionStrLenCP* expr) final {} + void visit(ExpressionSubtract* expr) final {} + void visit(ExpressionSwitch* expr) final {} + void visit(ExpressionToLower* expr) final {} + void visit(ExpressionToUpper* expr) final {} + void visit(ExpressionTrim* expr) final {} + void visit(ExpressionTrunc* expr) final {} + void visit(ExpressionType* expr) final {} + void visit(ExpressionZip* expr) final {} + void visit(ExpressionConvert* expr) final {} + void visit(ExpressionRegexFind* expr) final {} + void visit(ExpressionRegexFindAll* expr) final {} + void visit(ExpressionRegexMatch* expr) final {} + void visit(ExpressionCosine* expr) final {} + void visit(ExpressionSine* expr) final {} + void visit(ExpressionTangent* expr) final {} + void visit(ExpressionArcCosine* expr) final {} + void visit(ExpressionArcSine* expr) final {} + void visit(ExpressionArcTangent* expr) final {} + void visit(ExpressionArcTangent2* expr) final {} + void visit(ExpressionHyperbolicArcTangent* expr) final {} + void visit(ExpressionHyperbolicArcCosine* expr) final {} + void visit(ExpressionHyperbolicArcSine* expr) final {} + void visit(ExpressionHyperbolicTangent* expr) final {} + void visit(ExpressionHyperbolicCosine* expr) final {} + void visit(ExpressionHyperbolicSine* expr) final {} + void visit(ExpressionDegreesToRadians* expr) final {} + void visit(ExpressionRadiansToDegrees* expr) final {} + void visit(ExpressionDayOfMonth* expr) final {} + void visit(ExpressionDayOfWeek* expr) final {} + void visit(ExpressionDayOfYear* expr) final {} + void visit(ExpressionHour* expr) final {} + void visit(ExpressionMillisecond* expr) final {} + void visit(ExpressionMinute* expr) final {} + void visit(ExpressionMonth* expr) final {} + void visit(ExpressionSecond* expr) final {} + void visit(ExpressionWeek* expr) final {} + void visit(ExpressionIsoWeekYear* expr) final {} + void visit(ExpressionIsoDayOfWeek* expr) final {} + void visit(ExpressionIsoWeek* expr) final {} + void visit(ExpressionYear* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorAvg>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorMax>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorMin>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorStdDevPop>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorStdDevSamp>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorSum>* expr) final {} + void visit(ExpressionFromAccumulator<AccumulatorMergeObjects>* expr) final {} + void visit(ExpressionTests::Testable* expr) final {} + void visit(ExpressionInternalJsEmit* expr) final {} + void visit(ExpressionInternalFindSlice* expr) final {} + void visit(ExpressionInternalFindPositional* expr) final {} + void visit(ExpressionInternalFindElemMatch* expr) final {} + void visit(ExpressionFunction* expr) final {} + void visit(ExpressionInternalRemoveFieldTombstones* expr) final {} + void visit(ExpressionRandom* expr) final {} + +private: + void visitMultiBranchLogicExpression(Expression* expr, sbe::EPrimBinary::Op logicOp) { + // The infix visitor should only visit expressions with more than one child. + invariant(expr->getChildren().size() >= 2); + invariant(logicOp == sbe::EPrimBinary::logicOr || logicOp == sbe::EPrimBinary::logicAnd); + + auto frameId = _context->frameIdGenerator->generate(); + auto branchExpr = generateExpressionForLogicBranch(sbe::EVariable{frameId, 0}); + std::unique_ptr<sbe::EExpression> shortCircuitCondition; + if (logicOp == sbe::EPrimBinary::logicAnd) { + // The filter should take the short circuit path when the branch resolves to _false_, so + // we invert the filter condition. + shortCircuitCondition = sbe::makeE<sbe::ELocalBind>( + frameId, + sbe::makeEs(_context->popExpr()), + sbe::makeE<sbe::EPrimUnary>(sbe::EPrimUnary::logicNot, std::move(branchExpr))); + } else { + // For $or, keep the filter condition as is; the filter will take the short circuit path + // when the branch resolves to true. + shortCircuitCondition = sbe::makeE<sbe::ELocalBind>( + frameId, sbe::makeEs(_context->popExpr()), std::move(branchExpr)); + } + + auto branchStage = sbe::makeS<sbe::FilterStage<false>>(std::move(_context->traverseStage), + std::move(shortCircuitCondition)); + + auto& currentFrameStack = _context->logicalExpressionEvalFrameStack.top(); + currentFrameStack.branches.emplace_back( + std::make_pair(currentFrameStack.nextBranchResultSlot, std::move(branchStage))); + + if (currentFrameStack.branches.size() < (expr->getChildren().size() - 1)) { + _context->prepareToTranslateShortCircuitingBranch( + logicOp, _context->slotIdGenerator->generate()); + } else { + // We have already translated all but one of the branches, meaning the next branch we + // translate will be the final one and does not need an short-circuit logic. + _context->prepareToTranslateConcludingLogicalBranch(); + } + } + + ExpressionVisitorContext* _context; +}; + +class ExpressionPostVisitor final : public ExpressionVisitor { +public: + ExpressionPostVisitor(ExpressionVisitorContext* context) : _context{context} {} + + void visit(ExpressionConstant* expr) final { + auto [tag, val] = convertFrom(expr->getValue()); + _context->pushExpr(sbe::makeE<sbe::EConstant>(tag, val)); + } + + void visit(ExpressionAbs* expr) final { + auto frameId = _context->frameIdGenerator->generate(); + auto binds = sbe::makeEs(_context->popExpr()); + sbe::EVariable inputRef(frameId, 0); + + auto checkNullOrEmpty = sbe::makeE<sbe::EPrimBinary>( + sbe::EPrimBinary::logicOr, + sbe::makeE<sbe::EPrimUnary>( + sbe::EPrimUnary::logicNot, + sbe::makeE<sbe::EFunction>("exists", sbe::makeEs(inputRef.clone()))), + sbe::makeE<sbe::EFunction>("isNull", sbe::makeEs(inputRef.clone()))); + + auto absExpr = sbe::makeE<sbe::EIf>( + std::move(checkNullOrEmpty), + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Null, 0), + sbe::makeE<sbe::EIf>( + sbe::makeE<sbe::EFunction>("isNumber", sbe::makeEs(inputRef.clone())), + sbe::makeE<sbe::EFunction>("abs", sbe::makeEs(inputRef.clone())), + sbe::makeE<sbe::EFail>(ErrorCodes::Error{4822870}, + "$abs only supports numeric types, not string"))); + + _context->pushExpr( + sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(absExpr))); + } + void visit(ExpressionAdd* expr) final { + _context->ensureArity(2); + auto rhs = _context->popExpr(); + auto lhs = _context->popExpr(); + _context->pushExpr( + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::add, std::move(lhs), std::move(rhs))); + } + void visit(ExpressionAllElementsTrue* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionAnd* expr) final { + visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicAnd); + } + void visit(ExpressionAnyElementTrue* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionArray* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionArrayElemAt* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFirst* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionLast* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionObjectToArray* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionArrayToObject* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionBsonSize* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionCeil* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionCoerceToBool* expr) final { + unsupportedExpression("$coerceToBool"); + } + void visit(ExpressionCompare* expr) final { + _context->ensureArity(2); + std::vector<std::unique_ptr<sbe::EExpression>> operands(2); + for (auto it = operands.rbegin(); it != operands.rend(); ++it) { + *it = _context->popExpr(); + } + + auto frameId = _context->frameIdGenerator->generate(); + sbe::EVariable lhsRef(frameId, 0); + sbe::EVariable rhsRef(frameId, 1); + + auto comparisonOperator = [expr]() { + switch (expr->getOp()) { + case ExpressionCompare::CmpOp::EQ: + return sbe::EPrimBinary::eq; + case ExpressionCompare::CmpOp::NE: + return sbe::EPrimBinary::neq; + case ExpressionCompare::CmpOp::GT: + return sbe::EPrimBinary::greater; + case ExpressionCompare::CmpOp::GTE: + return sbe::EPrimBinary::greaterEq; + case ExpressionCompare::CmpOp::LT: + return sbe::EPrimBinary::less; + case ExpressionCompare::CmpOp::LTE: + return sbe::EPrimBinary::lessEq; + case ExpressionCompare::CmpOp::CMP: + return sbe::EPrimBinary::cmp3w; + } + MONGO_UNREACHABLE; + }(); + + // We use the "cmp3e" primitive for every comparison, because it "type brackets" its + // comparisons (for example, a number will always compare as less than a string). The other + // comparison primitives are designed for comparing values of the same type. + auto cmp3w = + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::cmp3w, lhsRef.clone(), rhsRef.clone()); + auto cmp = (comparisonOperator == sbe::EPrimBinary::cmp3w) + ? std::move(cmp3w) + : sbe::makeE<sbe::EPrimBinary>( + comparisonOperator, + std::move(cmp3w), + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt32, 0)); + + // If either operand evaluates to "Nothing," then the entire operation expressed by 'cmp' + // will also evaluate to "Nothing." MQL comparisons, however, treat "Nothing" as if it is a + // value that is less than everything other than MinKey. (Notably, two expressions that + // evaluate to "Nothing" are considered equal to each other.) + auto nothingFallbackCmp = sbe::makeE<sbe::EPrimBinary>( + comparisonOperator, + sbe::makeE<sbe::EFunction>("exists", sbe::makeEs(lhsRef.clone())), + sbe::makeE<sbe::EFunction>("exists", sbe::makeEs(rhsRef.clone()))); + + auto cmpWithFallback = sbe::makeE<sbe::EFunction>( + "fillEmpty", sbe::makeEs(std::move(cmp), std::move(nothingFallbackCmp))); + + _context->pushExpr( + sbe::makeE<sbe::ELocalBind>(frameId, std::move(operands), std::move(cmpWithFallback))); + } + void visit(ExpressionConcat* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionConcatArrays* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionCond* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionDateFromString* expr) final { + unsupportedExpression("$dateFromString"); + } + void visit(ExpressionDateFromParts* expr) final { + unsupportedExpression("$dateFromString"); + } + void visit(ExpressionDateToParts* expr) final { + unsupportedExpression("$dateFromString"); + } + void visit(ExpressionDateToString* expr) final { + unsupportedExpression("$dateFromString"); + } + void visit(ExpressionDivide* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionExp* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFieldPath* expr) final { + if (expr->getVariableId() == Variables::kRemoveId) { + // The case of $$REMOVE. Note that MQL allows a path in this situation (e.g., + // "$$REMOVE.foo.bar") but ignores it. + _context->pushExpr(sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Nothing, 0)); + return; + } + + sbe::value::SlotId slotId; + if (expr->isRootFieldPath()) { + slotId = _context->rootSlot; + } else { + auto it = _context->environment.find(expr->getVariableId()); + invariant(it != _context->environment.end()); + slotId = it->second; + } + + if (expr->getFieldPath().getPathLength() == 1) { + // A solo variable reference (e.g.: "$$ROOT" or "$$myvar") that doesn't need any + // traversal. + _context->pushExpr(sbe::makeE<sbe::EVariable>(slotId)); + return; + } + + // Dereference a dotted path, which may contain arrays requiring implicit traversal. + const bool expectsDocumentInputOnly = slotId == _context->rootSlot; + auto [outputSlot, stage] = generateTraverse(std::move(_context->traverseStage), + slotId, + expectsDocumentInputOnly, + expr->getFieldPathWithoutCurrentPrefix(), + _context->slotIdGenerator); + _context->pushExpr(sbe::makeE<sbe::EVariable>(outputSlot), std::move(stage)); + _context->relevantSlots->push_back(outputSlot); + } + void visit(ExpressionFilter* expr) final { + unsupportedExpression("$filter"); + } + void visit(ExpressionFloor* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionIfNull* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionIn* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionIndexOfArray* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionIndexOfBytes* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionIndexOfCP* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionIsNumber* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionLet* expr) final { + // The evaluated result of the $let is the evaluated result of its "in" field, which is + // already on top of the stack. The "infix" visitor has already popped the variable + // initializers off the expression stack. + _context->ensureArity(1); + + // We should have bound all the variables from this $let expression. + invariant(!_context->varsFrameStack.empty()); + auto& currentFrame = _context->varsFrameStack.top(); + invariant(currentFrame.variablesToBind.empty()); + + // Pop the lexical frame for this $let and remove all its bindings, which are now out of + // scope. + auto it = _context->environment.begin(); + while (it != _context->environment.end()) { + if (currentFrame.boundVariables.count(it->first)) { + it = _context->environment.erase(it); + } else { + ++it; + } + } + _context->varsFrameStack.pop(); + + // Note that there is no need to remove SlotId bindings from the the _context's environment. + // The AST parser already enforces scope rules. + } + void visit(ExpressionLn* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionLog* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionLog10* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionMap* expr) final { + unsupportedExpression("$map"); + } + void visit(ExpressionMeta* expr) final { + unsupportedExpression("$meta"); + } + void visit(ExpressionMod* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionMultiply* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionNot* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionObject* expr) final { + unsupportedExpression("$object"); + } + void visit(ExpressionOr* expr) final { + visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicOr); + } + void visit(ExpressionPow* expr) final { + unsupportedExpression("$pow"); + } + void visit(ExpressionRange* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionReduce* expr) final { + unsupportedExpression("$reduce"); + } + void visit(ExpressionReplaceOne* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionReplaceAll* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSetDifference* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSetEquals* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSetIntersection* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSetIsSubset* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSetUnion* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSize* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionReverseArray* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSlice* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionIsArray* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionRound* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSplit* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSqrt* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionStrcasecmp* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSubstrBytes* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSubstrCP* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionStrLenBytes* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionBinarySize* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionStrLenCP* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSubtract* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionSwitch* expr) final { + unsupportedExpression("$switch"); + } + void visit(ExpressionToLower* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionToUpper* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionTrim* expr) final { + unsupportedExpression("$trim"); + } + void visit(ExpressionTrunc* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionType* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionZip* expr) final { + unsupportedExpression("$zip"); + } + void visit(ExpressionConvert* expr) final { + unsupportedExpression("$convert"); + } + void visit(ExpressionRegexFind* expr) final { + unsupportedExpression("$regexFind"); + } + void visit(ExpressionRegexFindAll* expr) final { + unsupportedExpression("$regexFind"); + } + void visit(ExpressionRegexMatch* expr) final { + unsupportedExpression("$regexFind"); + } + void visit(ExpressionCosine* expr) final { + unsupportedExpression("$cos"); + } + void visit(ExpressionSine* expr) final { + unsupportedExpression("$sin"); + } + void visit(ExpressionTangent* expr) final { + unsupportedExpression("$tan"); + } + void visit(ExpressionArcCosine* expr) final { + unsupportedExpression("$acos"); + } + void visit(ExpressionArcSine* expr) final { + unsupportedExpression("$asin"); + } + void visit(ExpressionArcTangent* expr) final { + unsupportedExpression("$atan"); + } + void visit(ExpressionArcTangent2* expr) final { + unsupportedExpression("$atan2"); + } + void visit(ExpressionHyperbolicArcTangent* expr) final { + unsupportedExpression("$atanh"); + } + void visit(ExpressionHyperbolicArcCosine* expr) final { + unsupportedExpression("$acosh"); + } + void visit(ExpressionHyperbolicArcSine* expr) final { + unsupportedExpression("$asinh"); + } + void visit(ExpressionHyperbolicTangent* expr) final { + unsupportedExpression("$tanh"); + } + void visit(ExpressionHyperbolicCosine* expr) final { + unsupportedExpression("$cosh"); + } + void visit(ExpressionHyperbolicSine* expr) final { + unsupportedExpression("$sinh"); + } + void visit(ExpressionDegreesToRadians* expr) final { + unsupportedExpression("$degreesToRadians"); + } + void visit(ExpressionRadiansToDegrees* expr) final { + unsupportedExpression("$radiansToDegrees"); + } + void visit(ExpressionDayOfMonth* expr) final { + unsupportedExpression("$dayOfMonth"); + } + void visit(ExpressionDayOfWeek* expr) final { + unsupportedExpression("$dayOfWeek"); + } + void visit(ExpressionDayOfYear* expr) final { + unsupportedExpression("$dayOfYear"); + } + void visit(ExpressionHour* expr) final { + unsupportedExpression("$hour"); + } + void visit(ExpressionMillisecond* expr) final { + unsupportedExpression("$millisecond"); + } + void visit(ExpressionMinute* expr) final { + unsupportedExpression("$minute"); + } + void visit(ExpressionMonth* expr) final { + unsupportedExpression("$month"); + } + void visit(ExpressionSecond* expr) final { + unsupportedExpression("$second"); + } + void visit(ExpressionWeek* expr) final { + unsupportedExpression("$week"); + } + void visit(ExpressionIsoWeekYear* expr) final { + unsupportedExpression("$isoWeekYear"); + } + void visit(ExpressionIsoDayOfWeek* expr) final { + unsupportedExpression("$isoDayOfWeek"); + } + void visit(ExpressionIsoWeek* expr) final { + unsupportedExpression("$isoWeek"); + } + void visit(ExpressionYear* expr) final { + unsupportedExpression("$year"); + } + void visit(ExpressionFromAccumulator<AccumulatorAvg>* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFromAccumulator<AccumulatorMax>* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFromAccumulator<AccumulatorMin>* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFromAccumulator<AccumulatorStdDevPop>* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFromAccumulator<AccumulatorStdDevSamp>* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFromAccumulator<AccumulatorSum>* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionFromAccumulator<AccumulatorMergeObjects>* expr) final { + unsupportedExpression(expr->getOpName()); + } + void visit(ExpressionTests::Testable* expr) final { + unsupportedExpression("$test"); + } + void visit(ExpressionInternalJsEmit* expr) final { + unsupportedExpression("$internalJsEmit"); + } + void visit(ExpressionInternalFindSlice* expr) final { + unsupportedExpression("$internalFindSlice"); + } + void visit(ExpressionInternalFindPositional* expr) final { + unsupportedExpression("$internalFindPositional"); + } + void visit(ExpressionInternalFindElemMatch* expr) final { + unsupportedExpression("$internalFindElemMatch"); + } + void visit(ExpressionFunction* expr) final { + unsupportedExpression("$function"); + } + void visit(ExpressionInternalRemoveFieldTombstones* expr) final { + unsupportedExpression("$internalRemoveFieldTombstones"); + } + + void visit(ExpressionRandom* expr) final { + unsupportedExpression(expr->getOpName()); + } + +private: + /** + * Shared logic for $and, $or. Converts each child into an EExpression that evaluates to Boolean + * true or false, based on MQL rules for $and and $or branches, and then chains the branches + * together using binary and/or EExpressions so that the result has MQL's short-circuit + * semantics. + */ + void visitMultiBranchLogicExpression(Expression* expr, sbe::EPrimBinary::Op logicOp) { + invariant(logicOp == sbe::EPrimBinary::logicAnd || logicOp == sbe::EPrimBinary::logicOr); + + if (expr->getChildren().size() == 0) { + // Empty $and and $or always evaluate to their logical operator's identity value: true + // and false, respectively. + bool logicIdentityVal = (logicOp == sbe::EPrimBinary::logicAnd); + _context->pushExpr(sbe::makeE<sbe::EConstant>( + sbe::value::TypeTags::Boolean, sbe::value::bitcastFrom(logicIdentityVal))); + return; + } else if (expr->getChildren().size() == 1) { + // No need for short circuiting logic in a singleton $and/$or. Just execute the branch + // and return its result as a bool. + auto frameId = _context->frameIdGenerator->generate(); + _context->pushExpr(sbe::makeE<sbe::ELocalBind>( + frameId, + sbe::makeEs(_context->popExpr()), + generateExpressionForLogicBranch(sbe::EVariable{frameId, 0}))); + + return; + } + + auto& LogicalExpressionEvalFrame = _context->logicalExpressionEvalFrameStack.top(); + + // The last branch works differently from the others. It just uses a project stage to + // produce a true or false value for the branch result. + auto frameId = _context->frameIdGenerator->generate(); + auto lastBranchExpr = sbe::makeE<sbe::ELocalBind>( + frameId, + sbe::makeEs(_context->popExpr()), + generateExpressionForLogicBranch(sbe::EVariable{frameId, 0})); + auto lastBranchResultSlot = _context->slotIdGenerator->generate(); + auto lastBranch = sbe::makeProjectStage( + std::move(_context->traverseStage), lastBranchResultSlot, std::move(lastBranchExpr)); + LogicalExpressionEvalFrame.branches.emplace_back( + std::make_pair(lastBranchResultSlot, std::move(lastBranch))); + + std::vector<sbe::value::SlotVector> branchSlots; + std::vector<std::unique_ptr<sbe::PlanStage>> branchStages; + for (auto&& [slot, stage] : LogicalExpressionEvalFrame.branches) { + branchSlots.push_back(sbe::makeSV(slot)); + branchStages.push_back(std::move(stage)); + } + + auto branchResultSlot = _context->slotIdGenerator->generate(); + auto unionOfBranches = sbe::makeS<sbe::UnionStage>( + std::move(branchStages), std::move(branchSlots), sbe::makeSV(branchResultSlot)); + + // Restore 'relevantSlots' to the way it was before we started translating the logic + // operator. + *_context->relevantSlots = std::move(LogicalExpressionEvalFrame.savedRelevantSlots); + + // Get a list of slots that are used by $let expressions. These slots need to be available + // to the inner side of the LoopJoinStage, in case any of the branches want to reference one + // of the variables bound by the $let. + sbe::value::SlotVector letBindings; + for (auto&& [_, slot] : _context->environment) { + letBindings.push_back(slot); + } + + // The LoopJoinStage we are creating here will not expose any of the slots from its outer + // side except for the ones we explicity ask for. For that reason, we maintain the + // 'relevantSlots' list of slots that may still be referenced above this stage. All of the + // slots in 'letBindings' are relevant by this definition, but we track them separately, + // which is why we need to add them in now. + auto relevantSlotsWithLetBindings(*_context->relevantSlots); + relevantSlotsWithLetBindings.insert( + relevantSlotsWithLetBindings.end(), letBindings.begin(), letBindings.end()); + + // Put the union into a nested loop. The inner side of the nested loop will execute exactly + // once, trying each branch of the union until one of them short circuits or until it + // reaches the end. This process also restores the old 'traverseStage' value from before we + // started translating the logic operator, by placing it below the new nested loop stage. + auto stage = sbe::makeS<sbe::LoopJoinStage>( + std::move(LogicalExpressionEvalFrame.savedTraverseStage), + sbe::makeS<sbe::LimitSkipStage>(std::move(unionOfBranches), 1, boost::none), + std::move(relevantSlotsWithLetBindings), + std::move(letBindings), + nullptr /* predicate */); + + // We've already restored all necessary state from the top 'logicalExpressionEvalFrameStack' + // entry, so we are done with it. + _context->logicalExpressionEvalFrameStack.pop(); + + // The final true/false result of the logic operator is stored in the 'branchResultSlot' + // slot. + _context->relevantSlots->push_back(branchResultSlot); + _context->pushExpr(sbe::makeE<sbe::EVariable>(branchResultSlot), std::move(stage)); + } + + void unsupportedExpression(const char* op) const { + uasserted(ErrorCodes::InternalErrorNotSupported, + str::stream() << "Expression is not supported in SBE: " << op); + } + + ExpressionVisitorContext* _context; +}; // namespace + +class ExpressionWalker final { +public: + ExpressionWalker(ExpressionVisitor* preVisitor, + ExpressionVisitor* inVisitor, + ExpressionVisitor* postVisitor) + : _preVisitor{preVisitor}, _inVisitor{inVisitor}, _postVisitor{postVisitor} {} + + void preVisit(Expression* expr) { + expr->acceptVisitor(_preVisitor); + } + + void inVisit(long long count, Expression* expr) { + expr->acceptVisitor(_inVisitor); + } + + void postVisit(Expression* expr) { + expr->acceptVisitor(_postVisitor); + } + +private: + ExpressionVisitor* _preVisitor; + ExpressionVisitor* _inVisitor; + ExpressionVisitor* _postVisitor; +}; +} // namespace + +std::tuple<sbe::value::SlotId, std::unique_ptr<sbe::EExpression>, std::unique_ptr<sbe::PlanStage>> +generateExpression(Expression* expr, + std::unique_ptr<sbe::PlanStage> stage, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::FrameIdGenerator* frameIdGenerator, + sbe::value::SlotId rootSlot, + sbe::value::SlotVector* relevantSlots) { + auto tempRelevantSlots = sbe::makeSV(rootSlot); + relevantSlots = relevantSlots ? relevantSlots : &tempRelevantSlots; + + ExpressionVisitorContext context( + std::move(stage), slotIdGenerator, frameIdGenerator, rootSlot, relevantSlots); + ExpressionPreVisitor preVisitor{&context}; + ExpressionInVisitor inVisitor{&context}; + ExpressionPostVisitor postVisitor{&context}; + ExpressionWalker walker{&preVisitor, &inVisitor, &postVisitor}; + expression_walker::walk(&walker, expr); + return context.done(); +} +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_expression.h b/src/mongo/db/query/sbe_stage_builder_expression.h new file mode 100644 index 00000000000..53c935b9455 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_expression.h @@ -0,0 +1,58 @@ +/** + * 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" +#include "mongo/db/pipeline/expression.h" + +namespace mongo::stage_builder { +/** + * Translates an input Expression into an SBE EExpression, along with a chain of PlanStages whose + * output will be necessary to evaluate the EExpression. The 'stage' input will be attached to the + * end of the resulting chain of PlanStages. + * + * Note that any slot whose value must be visible to the parent of the PlanStage output by this + * function should be included in the 'relevantSlots' list. Some stages (notably LoopJoin) do not + * forward all of the slots visible to them to their parents; they need an explicit list of which + * slots to forward. + * + * The 'relevantSlots' is an input/output parameter. Execution of this function may add additional + * relevant slots to thie list. + */ +std::tuple<sbe::value::SlotId, std::unique_ptr<sbe::EExpression>, std::unique_ptr<sbe::PlanStage>> +generateExpression(Expression* expr, + std::unique_ptr<sbe::PlanStage> stage, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::FrameIdGenerator* frameIdGenerator, + sbe::value::SlotId inputVar, + sbe::value::SlotVector* relevantSlots = nullptr); + +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_filter.cpp b/src/mongo/db/query/sbe_stage_builder_filter.cpp new file mode 100644 index 00000000000..8d5fc79b715 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_filter.cpp @@ -0,0 +1,726 @@ +/** + * 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/query/sbe_stage_builder_filter.h" + +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/stages/traverse.h" +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/matcher/expression_always_boolean.h" +#include "mongo/db/matcher/expression_array.h" +#include "mongo/db/matcher/expression_expr.h" +#include "mongo/db/matcher/expression_geo.h" +#include "mongo/db/matcher/expression_internal_expr_eq.h" +#include "mongo/db/matcher/expression_leaf.h" +#include "mongo/db/matcher/expression_text.h" +#include "mongo/db/matcher/expression_text_noop.h" +#include "mongo/db/matcher/expression_tree.h" +#include "mongo/db/matcher/expression_type.h" +#include "mongo/db/matcher/expression_visitor.h" +#include "mongo/db/matcher/expression_where.h" +#include "mongo/db/matcher/expression_where_noop.h" +#include "mongo/db/matcher/schema/expression_internal_schema_all_elem_match_from_index.h" +#include "mongo/db/matcher/schema/expression_internal_schema_allowed_properties.h" +#include "mongo/db/matcher/schema/expression_internal_schema_cond.h" +#include "mongo/db/matcher/schema/expression_internal_schema_eq.h" +#include "mongo/db/matcher/schema/expression_internal_schema_fmod.h" +#include "mongo/db/matcher/schema/expression_internal_schema_match_array_index.h" +#include "mongo/db/matcher/schema/expression_internal_schema_max_items.h" +#include "mongo/db/matcher/schema/expression_internal_schema_max_length.h" +#include "mongo/db/matcher/schema/expression_internal_schema_max_properties.h" +#include "mongo/db/matcher/schema/expression_internal_schema_min_items.h" +#include "mongo/db/matcher/schema/expression_internal_schema_min_length.h" +#include "mongo/db/matcher/schema/expression_internal_schema_min_properties.h" +#include "mongo/db/matcher/schema/expression_internal_schema_object_match.h" +#include "mongo/db/matcher/schema/expression_internal_schema_root_doc_eq.h" +#include "mongo/db/matcher/schema/expression_internal_schema_unique_items.h" +#include "mongo/db/matcher/schema/expression_internal_schema_xor.h" +#include "mongo/logv2/log.h" +#include "mongo/util/str.h" + +namespace mongo::stage_builder { +namespace { +/** + * The various flavors of PathMatchExpressions require the same skeleton of traverse operators in + * order to perform implicit path traversal, but may translate differently to an SBE expression that + * actually applies the predicate against an individual array element. + * + * A function of this type can be called to generate an EExpression which applies a predicate to the + * value found in 'inputSlot'. + */ +using MakePredicateEExprFn = + std::function<std::unique_ptr<sbe::EExpression>(sbe::value::SlotId inputSlot)>; + +/** + * A struct for storing context across calls to visit() methods in MatchExpressionVisitor's. + */ +struct MatchExpressionVisitorContext { + MatchExpressionVisitorContext(sbe::value::SlotIdGenerator* slotIdGenerator, + std::unique_ptr<sbe::PlanStage> inputStage, + sbe::value::SlotId inputVar) + : slotIdGenerator{slotIdGenerator}, inputStage{std::move(inputStage)}, inputVar{inputVar} {} + + std::unique_ptr<sbe::PlanStage> done() { + if (!predicateVars.empty()) { + invariant(predicateVars.size() == 1); + inputStage = sbe::makeS<sbe::FilterStage<false>>( + std::move(inputStage), sbe::makeE<sbe::EVariable>(predicateVars.top())); + predicateVars.pop(); + } + return std::move(inputStage); + } + + sbe::value::SlotIdGenerator* slotIdGenerator; + std::unique_ptr<sbe::PlanStage> inputStage; + std::stack<sbe::value::SlotId> predicateVars; + std::stack<std::pair<const MatchExpression*, size_t>> nestedLogicalExprs; + sbe::value::SlotId inputVar; +}; + +/** + * A match expression tree walker to be used with MatchExpression visitors in order to translate + * a MatchExpression tree into an SBE plane stage sub-tree which implements the filter. + */ +class MatchExpressionWalker final { +public: + MatchExpressionWalker(MatchExpressionConstVisitor* preVisitor, + MatchExpressionConstVisitor* inVisitor, + MatchExpressionConstVisitor* postVisitor) + : _preVisitor{preVisitor}, _inVisitor{inVisitor}, _postVisitor{postVisitor} {} + + void preVisit(const MatchExpression* expr) { + expr->acceptVisitor(_preVisitor); + } + + void postVisit(const MatchExpression* expr) { + expr->acceptVisitor(_postVisitor); + } + + void inVisit(long count, const MatchExpression* expr) { + expr->acceptVisitor(_inVisitor); + } + +private: + MatchExpressionConstVisitor* _preVisitor; + MatchExpressionConstVisitor* _inVisitor; + MatchExpressionConstVisitor* _postVisitor; +}; + +/** + * A helper function to generate a path traversal plan stage at the given nested 'level' of the + * traversal path. For example, for a dotted path expression {'a.b': 2}, the traversal sub-tree will + * look like this: + * + * traverse + * traversePredicateVar // the global traversal result + * elemPredicateVar1 // the result coming from the 'in' branch + * fieldVar1 // field 'a' projected in the 'from' branch, this is the field we will be + * // traversing + * {traversePredicateVar || elemPredicateVar1} // the folding expression - combining + * // results for each element + * {traversePredicateVar} // final (early out) expression - when we hit the 'true' value, + * // we don't have to traverse the whole array + * in + * project [elemPredicateVar1 = traversePredicateVar] + * traverse // nested traversal + * traversePredicateVar // the global traversal result + * elemPredicateVar2 // the result coming from the 'in' branch + * fieldVar2 // field 'b' projected in the 'from' branch, this is the field we will be + * // traversing + * {traversePredicateVar || elemPredicateVar2} // the folding expression + * {traversePredicateVar} // final (early out) expression + * in + * project [elemPredicateVar2 = fieldVar2==2] // compare the field 'b' to 2 and store + * // the bool result in elemPredicateVar2 + * limit 1 + * coscan + * from + * project [fieldVar2=getField(fieldVar1, 'b')] // project field 'b' from the document + * // bound to 'fieldVar1', which is + * // field 'a' + * limit 1 + * coscan + * from + * project [fieldVar1=getField(inputVar, 'a')] // project field 'a' from the document bound + * // to 'inputVar' + * <inputStage> // e.g., COLLSCAN + */ +std::unique_ptr<sbe::PlanStage> generateTraverseHelper(MatchExpressionVisitorContext* context, + std::unique_ptr<sbe::PlanStage> inputStage, + sbe::value::SlotId inputVar, + const PathMatchExpression* expr, + MakePredicateEExprFn makeEExprCallback, + size_t level) { + using namespace std::literals; + + FieldPath path{expr->path()}; + invariant(level < path.getPathLength()); + + // The global traversal result. + const auto& traversePredicateVar = context->predicateVars.top(); + // The field we will be traversing at the current nested level. + auto fieldVar{context->slotIdGenerator->generate()}; + // The result coming from the 'in' branch of the traverse plan stage. + auto elemPredicateVar{context->slotIdGenerator->generate()}; + + // Generate the projection stage to read a sub-field at the current nested level and bind it + // to 'fieldVar'. + inputStage = sbe::makeProjectStage( + std::move(inputStage), + fieldVar, + sbe::makeE<sbe::EFunction>( + "getField"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(inputVar), sbe::makeE<sbe::EConstant>([&]() { + auto fieldName = path.getFieldName(level); + return std::string_view{fieldName.rawData(), fieldName.size()}; + }())))); + + std::unique_ptr<sbe::PlanStage> innerBranch; + if (level == path.getPathLength() - 1u) { + innerBranch = sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + elemPredicateVar, + makeEExprCallback(fieldVar)); + } else { + // Generate nested traversal. + innerBranch = sbe::makeProjectStage( + generateTraverseHelper( + context, + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + fieldVar, + expr, + makeEExprCallback, + level + 1), + elemPredicateVar, + sbe::makeE<sbe::EVariable>(traversePredicateVar)); + } + + // The final traverse stage for the current nested level. + return sbe::makeS<sbe::TraverseStage>( + std::move(inputStage), + std::move(innerBranch), + fieldVar, + traversePredicateVar, + elemPredicateVar, + sbe::makeSV(), + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr, + sbe::makeE<sbe::EVariable>(traversePredicateVar), + sbe::makeE<sbe::EVariable>(elemPredicateVar)), + sbe::makeE<sbe::EVariable>(traversePredicateVar), + 2); +} + +/** + * For the given PathMatchExpression 'expr', generates a path traversal SBE plan stage sub-tree + * implementing the expression. Generates a sequence of nested traverse operators in order to + * perform nested array traversal, and then calls 'makeEExprCallback' in order to generate an SBE + * expression responsible for applying the predicate to individual array elements. + */ +void generateTraverse(MatchExpressionVisitorContext* context, + const PathMatchExpression* expr, + MakePredicateEExprFn makeEExprCallback) { + context->predicateVars.push(context->slotIdGenerator->generate()); + context->inputStage = generateTraverseHelper(context, + std::move(context->inputStage), + context->inputVar, + expr, + std::move(makeEExprCallback), + 0); + + // If this comparison expression is a branch of a logical $and expression, but not the last + // one, inject a filter stage to bail out early from the $and predicate without the need to + // evaluate all branches. If this is the last branch of the $and expression, or if it's not + // within a logical expression at all, just keep the predicate var on the top on the stack + // and let the parent expression process it. + if (!context->nestedLogicalExprs.empty() && context->nestedLogicalExprs.top().second > 1 && + context->nestedLogicalExprs.top().first->matchType() == MatchExpression::AND) { + context->inputStage = sbe::makeS<sbe::FilterStage<false>>( + std::move(context->inputStage), + sbe::makeE<sbe::EVariable>(context->predicateVars.top())); + context->predicateVars.pop(); + } +} + +/** + * Generates a path traversal SBE plan stage sub-tree which implments the comparison match + * expression 'expr'. The comparison itself executes using the given 'binaryOp'. + */ +void generateTraverseForComparisonPredicate(MatchExpressionVisitorContext* context, + const ComparisonMatchExpression* expr, + sbe::EPrimBinary::Op binaryOp) { + auto makeEExprFn = [expr, binaryOp](sbe::value::SlotId inputSlot) { + const auto& rhs = expr->getData(); + auto [tagView, valView] = sbe::bson::convertFrom( + true, rhs.rawdata(), rhs.rawdata() + rhs.size(), rhs.fieldNameSize() - 1); + + // SBE EConstant assumes ownership of the value so we have to make a copy here. + auto [tag, val] = sbe::value::copyValue(tagView, valView); + + return sbe::makeE<sbe::EPrimBinary>( + binaryOp, sbe::makeE<sbe::EVariable>(inputSlot), sbe::makeE<sbe::EConstant>(tag, val)); + }; + generateTraverse(context, expr, std::move(makeEExprFn)); +} + +/** + * Generates an SBE plan stage sub-tree implementing a logical $or expression. + */ +void generateLogicalOr(MatchExpressionVisitorContext* context, const OrMatchExpression* expr) { + invariant(!context->predicateVars.empty()); + invariant(context->predicateVars.size() >= expr->numChildren()); + + auto filter = sbe::makeE<sbe::EVariable>(context->predicateVars.top()); + context->predicateVars.pop(); + + auto numOrBranches = expr->numChildren() - 1; + for (size_t childNum = 0; childNum < numOrBranches; ++childNum) { + filter = + sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr, + std::move(filter), + sbe::makeE<sbe::EVariable>(context->predicateVars.top())); + context->predicateVars.pop(); + } + + // If this $or expression is a branch of another $and expression, or is a top-level logical + // expression we can just inject a filter stage without propagating the result of the predicate + // evaluation to the parent expression, to form a sub-tree of stage->FILTER->stage->FILTER plan + // stages to support early exit for the $and branches. Otherwise, just project out the result + // of the predicate evaluation and let the parent expression handle it. + if (context->nestedLogicalExprs.empty() || + context->nestedLogicalExprs.top().first->matchType() == MatchExpression::AND) { + context->inputStage = + sbe::makeS<sbe::FilterStage<false>>(std::move(context->inputStage), std::move(filter)); + } else { + context->predicateVars.push(context->slotIdGenerator->generate()); + context->inputStage = sbe::makeProjectStage( + std::move(context->inputStage), context->predicateVars.top(), std::move(filter)); + } +} + +/** + * Generates an SBE plan stage sub-tree implementing a logical $and expression. + */ +void generateLogicalAnd(MatchExpressionVisitorContext* context, const AndMatchExpression* expr) { + auto filter = [&]() { + if (expr->numChildren() > 0) { + invariant(!context->predicateVars.empty()); + auto predicateVar = context->predicateVars.top(); + context->predicateVars.pop(); + return sbe::makeE<sbe::EVariable>(predicateVar); + } else { + return sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Boolean, 1); + } + }(); + + // If this $and expression is a branch of another $and expression, or is a top-level logical + // expression we can just inject a filter stage without propagating the result of the predicate + // evaluation to the parent expression, to form a sub-tree of stage->FILTER->stage->FILTER plan + // stages to support early exit for the $and branches. Otherwise, just project out the result + // of the predicate evaluation and let the parent expression handle it. + if (context->nestedLogicalExprs.empty() || + context->nestedLogicalExprs.top().first->matchType() == MatchExpression::AND) { + context->inputStage = + sbe::makeS<sbe::FilterStage<false>>(std::move(context->inputStage), std::move(filter)); + } else { + context->predicateVars.push(context->slotIdGenerator->generate()); + context->inputStage = sbe::makeProjectStage( + std::move(context->inputStage), context->predicateVars.top(), std::move(filter)); + } +} + +/** + * A match expression pre-visitor used for maintaining nested logical expressions while traversing + * the match expression tree. + */ +class MatchExpressionPreVisitor final : public MatchExpressionConstVisitor { +public: + MatchExpressionPreVisitor(MatchExpressionVisitorContext* context) : _context(context) {} + + void visit(const AlwaysFalseMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const AlwaysTrueMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const AndMatchExpression* expr) final { + _context->nestedLogicalExprs.push({expr, expr->numChildren()}); + } + void visit(const BitsAllClearMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const BitsAllSetMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const BitsAnyClearMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const BitsAnySetMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const ElemMatchObjectMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const ElemMatchValueMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const EqualityMatchExpression* expr) final {} + void visit(const ExistsMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const ExprMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const GTEMatchExpression* expr) final {} + void visit(const GTMatchExpression* expr) final {} + void visit(const GeoMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const GeoNearMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalExprEqMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaAllElemMatchFromIndexMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaAllowedPropertiesMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaBinDataEncryptedTypeExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaBinDataSubTypeExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaCondMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaEqMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaFmodMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaMatchArrayIndexMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaMaxItemsMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaMaxLengthMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaMaxPropertiesMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaMinItemsMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaMinLengthMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaMinPropertiesMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaObjectMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaRootDocEqMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaTypeExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaUniqueItemsMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const InternalSchemaXorMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const LTEMatchExpression* expr) final {} + void visit(const LTMatchExpression* expr) final {} + void visit(const ModMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const NorMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const NotMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const OrMatchExpression* expr) final { + _context->nestedLogicalExprs.push({expr, expr->numChildren()}); + } + void visit(const RegexMatchExpression* expr) final {} + void visit(const SizeMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const TextMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const TextNoOpMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const TwoDPtInAnnulusExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const TypeMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const WhereMatchExpression* expr) final { + unsupportedExpression(expr); + } + void visit(const WhereNoOpMatchExpression* expr) final { + unsupportedExpression(expr); + } + +private: + void unsupportedExpression(const MatchExpression* expr) const { + uasserted(4822878, + str::stream() << "Match expression is not supported in SBE: " + << expr->matchType()); + } + + MatchExpressionVisitorContext* _context; +}; + +/** + * A match expression post-visitor which does all the job to translate the match expression tree + * into an SBE plan stage sub-tree. + */ +class MatchExpressionPostVisitor final : public MatchExpressionConstVisitor { +public: + MatchExpressionPostVisitor(MatchExpressionVisitorContext* context) : _context(context) {} + + void visit(const AlwaysFalseMatchExpression* expr) final {} + void visit(const AlwaysTrueMatchExpression* expr) final {} + void visit(const AndMatchExpression* expr) final { + _context->nestedLogicalExprs.pop(); + generateLogicalAnd(_context, expr); + } + void visit(const BitsAllClearMatchExpression* expr) final {} + void visit(const BitsAllSetMatchExpression* expr) final {} + void visit(const BitsAnyClearMatchExpression* expr) final {} + void visit(const BitsAnySetMatchExpression* expr) final {} + void visit(const ElemMatchObjectMatchExpression* expr) final {} + void visit(const ElemMatchValueMatchExpression* expr) final {} + void visit(const EqualityMatchExpression* expr) final { + generateTraverseForComparisonPredicate(_context, expr, sbe::EPrimBinary::eq); + } + void visit(const ExistsMatchExpression* expr) final {} + void visit(const ExprMatchExpression* expr) final {} + void visit(const GTEMatchExpression* expr) final { + generateTraverseForComparisonPredicate(_context, expr, sbe::EPrimBinary::greaterEq); + } + void visit(const GTMatchExpression* expr) final { + generateTraverseForComparisonPredicate(_context, expr, sbe::EPrimBinary::greater); + } + void visit(const GeoMatchExpression* expr) final {} + void visit(const GeoNearMatchExpression* expr) final {} + void visit(const InMatchExpression* expr) final {} + void visit(const InternalExprEqMatchExpression* expr) final {} + void visit(const InternalSchemaAllElemMatchFromIndexMatchExpression* expr) final {} + void visit(const InternalSchemaAllowedPropertiesMatchExpression* expr) final {} + void visit(const InternalSchemaBinDataEncryptedTypeExpression* expr) final {} + void visit(const InternalSchemaBinDataSubTypeExpression* expr) final {} + void visit(const InternalSchemaCondMatchExpression* expr) final {} + void visit(const InternalSchemaEqMatchExpression* expr) final {} + void visit(const InternalSchemaFmodMatchExpression* expr) final {} + void visit(const InternalSchemaMatchArrayIndexMatchExpression* expr) final {} + void visit(const InternalSchemaMaxItemsMatchExpression* expr) final {} + void visit(const InternalSchemaMaxLengthMatchExpression* expr) final {} + void visit(const InternalSchemaMaxPropertiesMatchExpression* expr) final {} + void visit(const InternalSchemaMinItemsMatchExpression* expr) final {} + void visit(const InternalSchemaMinLengthMatchExpression* expr) final {} + void visit(const InternalSchemaMinPropertiesMatchExpression* expr) final {} + void visit(const InternalSchemaObjectMatchExpression* expr) final {} + void visit(const InternalSchemaRootDocEqMatchExpression* expr) final {} + void visit(const InternalSchemaTypeExpression* expr) final {} + void visit(const InternalSchemaUniqueItemsMatchExpression* expr) final {} + void visit(const InternalSchemaXorMatchExpression* expr) final {} + void visit(const LTEMatchExpression* expr) final { + generateTraverseForComparisonPredicate(_context, expr, sbe::EPrimBinary::lessEq); + } + void visit(const LTMatchExpression* expr) final { + generateTraverseForComparisonPredicate(_context, expr, sbe::EPrimBinary::less); + } + void visit(const ModMatchExpression* expr) final {} + void visit(const NorMatchExpression* expr) final {} + void visit(const NotMatchExpression* expr) final {} + void visit(const OrMatchExpression* expr) final { + _context->nestedLogicalExprs.pop(); + generateLogicalOr(_context, expr); + } + + void visit(const RegexMatchExpression* expr) final { + auto makeEExprFn = [expr](sbe::value::SlotId inputSlot) { + auto regex = RegexMatchExpression::makeRegex(expr->getString(), expr->getFlags()); + auto ownedRegexVal = sbe::value::bitcastFrom(regex.release()); + + // The "regexMatch" function returns Nothing when given any non-string input, so we need + // an explicit string check in the expression in order to capture the MQL semantics of + // regex returning false for non-strings. We generate the following expression: + // + // and + // +----------------+----------------+ + // isString regexMatch + // | +------------+----------+ + // var (inputSlot) constant (regex) var (inputSlot) + // + // TODO: In the future, this needs to account for the fact that the regex match + // expression matches strings, but also matches stored regexes. For example, + // {$match: {a: /foo/}} matches the document {a: /foo/} in addition to {a: "foobar"}. + return sbe::makeE<sbe::EPrimBinary>( + sbe::EPrimBinary::logicAnd, + sbe::makeE<sbe::EFunction>("isString", + sbe::makeEs(sbe::makeE<sbe::EVariable>(inputSlot))), + sbe::makeE<sbe::EFunction>( + "regexMatch", + sbe::makeEs( + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::pcreRegex, ownedRegexVal), + sbe::makeE<sbe::EVariable>(inputSlot)))); + }; + + generateTraverse(_context, expr, std::move(makeEExprFn)); + } + + void visit(const SizeMatchExpression* expr) final {} + void visit(const TextMatchExpression* expr) final {} + void visit(const TextNoOpMatchExpression* expr) final {} + void visit(const TwoDPtInAnnulusExpression* expr) final {} + void visit(const TypeMatchExpression* expr) final {} + void visit(const WhereMatchExpression* expr) final {} + void visit(const WhereNoOpMatchExpression* expr) final {} + +private: + MatchExpressionVisitorContext* _context; +}; + +/** + * A match expression in-visitor used for maintaining the counter of the processed child expressions + * of the nested logical expressions in the match expression tree being traversed. + */ +class MatchExpressionInVisitor final : public MatchExpressionConstVisitor { +public: + MatchExpressionInVisitor(MatchExpressionVisitorContext* context) : _context(context) {} + + void visit(const AlwaysFalseMatchExpression* expr) final {} + void visit(const AlwaysTrueMatchExpression* expr) final {} + void visit(const AndMatchExpression* expr) final { + invariant(_context->nestedLogicalExprs.top().first == expr); + _context->nestedLogicalExprs.top().second--; + } + void visit(const BitsAllClearMatchExpression* expr) final {} + void visit(const BitsAllSetMatchExpression* expr) final {} + void visit(const BitsAnyClearMatchExpression* expr) final {} + void visit(const BitsAnySetMatchExpression* expr) final {} + void visit(const ElemMatchObjectMatchExpression* expr) final {} + void visit(const ElemMatchValueMatchExpression* expr) final {} + void visit(const EqualityMatchExpression* expr) final {} + void visit(const ExistsMatchExpression* expr) final {} + void visit(const ExprMatchExpression* expr) final {} + void visit(const GTEMatchExpression* expr) final {} + void visit(const GTMatchExpression* expr) final {} + void visit(const GeoMatchExpression* expr) final {} + void visit(const GeoNearMatchExpression* expr) final {} + void visit(const InMatchExpression* expr) final {} + void visit(const InternalExprEqMatchExpression* expr) final {} + void visit(const InternalSchemaAllElemMatchFromIndexMatchExpression* expr) final {} + void visit(const InternalSchemaAllowedPropertiesMatchExpression* expr) final {} + void visit(const InternalSchemaBinDataEncryptedTypeExpression* expr) final {} + void visit(const InternalSchemaBinDataSubTypeExpression* expr) final {} + void visit(const InternalSchemaCondMatchExpression* expr) final {} + void visit(const InternalSchemaEqMatchExpression* expr) final {} + void visit(const InternalSchemaFmodMatchExpression* expr) final {} + void visit(const InternalSchemaMatchArrayIndexMatchExpression* expr) final {} + void visit(const InternalSchemaMaxItemsMatchExpression* expr) final {} + void visit(const InternalSchemaMaxLengthMatchExpression* expr) final {} + void visit(const InternalSchemaMaxPropertiesMatchExpression* expr) final {} + void visit(const InternalSchemaMinItemsMatchExpression* expr) final {} + void visit(const InternalSchemaMinLengthMatchExpression* expr) final {} + void visit(const InternalSchemaMinPropertiesMatchExpression* expr) final {} + void visit(const InternalSchemaObjectMatchExpression* expr) final {} + void visit(const InternalSchemaRootDocEqMatchExpression* expr) final {} + void visit(const InternalSchemaTypeExpression* expr) final {} + void visit(const InternalSchemaUniqueItemsMatchExpression* expr) final {} + void visit(const InternalSchemaXorMatchExpression* expr) final {} + void visit(const LTEMatchExpression* expr) final {} + void visit(const LTMatchExpression* expr) final {} + void visit(const ModMatchExpression* expr) final {} + void visit(const NorMatchExpression* expr) final {} + void visit(const NotMatchExpression* expr) final {} + void visit(const OrMatchExpression* expr) final { + invariant(_context->nestedLogicalExprs.top().first == expr); + _context->nestedLogicalExprs.top().second--; + } + void visit(const RegexMatchExpression* expr) final {} + void visit(const SizeMatchExpression* expr) final {} + void visit(const TextMatchExpression* expr) final {} + void visit(const TextNoOpMatchExpression* expr) final {} + void visit(const TwoDPtInAnnulusExpression* expr) final {} + void visit(const TypeMatchExpression* expr) final {} + void visit(const WhereMatchExpression* expr) final {} + void visit(const WhereNoOpMatchExpression* expr) final {} + +private: + MatchExpressionVisitorContext* _context; +}; +} // namespace + +std::unique_ptr<sbe::PlanStage> generateFilter(const MatchExpression* root, + std::unique_ptr<sbe::PlanStage> stage, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::SlotId inputVar) { + // The planner adds an $and expression without the operands if the query was empty. We can bail + // out early without generating the filter plan stage if this is the case. + if (root->matchType() == MatchExpression::AND && root->numChildren() == 0) { + return stage; + } + + MatchExpressionVisitorContext context{slotIdGenerator, std::move(stage), inputVar}; + MatchExpressionPreVisitor preVisitor{&context}; + MatchExpressionInVisitor inVisitor{&context}; + MatchExpressionPostVisitor postVisitor{&context}; + MatchExpressionWalker walker{&preVisitor, &inVisitor, &postVisitor}; + tree_walker::walk<true, MatchExpression>(root, &walker); + return context.done(); +} +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_filter.h b/src/mongo/db/query/sbe_stage_builder_filter.h new file mode 100644 index 00000000000..1bcf6283cb7 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_filter.h @@ -0,0 +1,46 @@ +/** + * 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/matcher/expression.h" + +namespace mongo::stage_builder { +/** + * Generates an SBE plan stage sub-tree implementing a filter expression represented by the 'root' + * expression. The 'stage' parameter defines an input stage to the generate SBE plan stage sub-tree. + * The 'inputVar' defines a variable to read the input document from. + */ +std::unique_ptr<sbe::PlanStage> generateFilter(const MatchExpression* root, + std::unique_ptr<sbe::PlanStage> stage, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::SlotId inputVar); + +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp new file mode 100644 index 00000000000..3b7085e9e6a --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp @@ -0,0 +1,670 @@ +/** + * 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/query/sbe_stage_builder_index_scan.h" + +#include "mongo/db/catalog/collection.h" +#include "mongo/db/exec/sbe/stages/check_bounds.h" +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/hash_agg.h" +#include "mongo/db/exec/sbe/stages/ix_scan.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/db/exec/sbe/stages/makeobj.h" +#include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/stages/spool.h" +#include "mongo/db/exec/sbe/stages/union.h" +#include "mongo/db/exec/sbe/stages/unwind.h" +#include "mongo/db/index/index_access_method.h" +#include "mongo/db/query/index_bounds_builder.h" +#include "mongo/db/query/util/make_data_structure.h" +#include "mongo/logv2/log.h" +#include "mongo/util/str.h" + +namespace mongo::stage_builder { +namespace { +/** + * Returns 'true' if the index bounds in 'intervalLists' can be represented as a number of intervals + * between low and high keys, which can be statically generated. Inclusivity of each bound is + * returned through the relevant '*KeyInclusive' parameter. Returns 'false' otherwise. + */ +bool canBeDecomposedIntoSingleIntervals(const std::vector<OrderedIntervalList>& intervalLists, + bool* lowKeyInclusive, + bool* highKeyInclusive) { + invariant(lowKeyInclusive); + invariant(highKeyInclusive); + + *lowKeyInclusive = true; + *highKeyInclusive = true; + + size_t listNum = 0; + + // First, we skip over point intervals. + for (; listNum < intervalLists.size(); ++listNum) { + if (!std::all_of(std::begin(intervalLists[listNum].intervals), + std::end(intervalLists[listNum].intervals), + [](auto&& interval) { return interval.isPoint(); })) { + break; + } + } + + // Bail out early if all our intervals are points. + if (listNum == intervalLists.size()) { + return true; + } + + // After point intervals we can have exactly one non-point interval. + if (intervalLists[listNum].intervals.size() != 1) { + return false; + } + + // Set the inclusivity from the non-point interval. + *lowKeyInclusive = intervalLists[listNum].intervals[0].startInclusive; + *highKeyInclusive = intervalLists[listNum].intervals[0].endInclusive; + + // And after the non-point interval we can have any number of "all values" intervals. + for (++listNum; listNum < intervalLists.size(); ++listNum) { + if (!(intervalLists[listNum].intervals.size() == 1 && + (intervalLists[listNum].intervals[0].isMinToMax() || + intervalLists[listNum].intervals[0].isMaxToMin()))) { + break; + } + } + + // If we've reached the end of the interval lists, then we can decompose a multi-interval index + // bounds into a number of single-interval bounds. + return listNum == intervalLists.size(); +} + +/** + * Decomposes multi-interval index bounds represented as 'intervalLists' into a number of + * single-interval bounds. Inclusivity of each bound is set through the relevant '*KeyInclusive' + * parameter. For example, if we've got an index {a: 1, b: 1, c: 1, d: 1} and would issue this + * query: + * + * {a: {$in: [1,2]}, b: {$in: [10,11]}, c: {$gte: 20}} + * + * Then the 'intervalLists' would contain the following multi-interval bounds: + * + * [ + * [ [1,1], [2,2] ], + * [ [10,10], [11,11] ], + * [ [20, Inf) ], + * [ [MinKey, MaxKey] + * ] + * + * And it'd be decomposed into the following single-intervals between low and high keys: + * + * {'':1, '':10, '':20, '':MinKey} -> {'':1, '':10, '':Inf, '':MaxKey} + * {'':1, '':11, '':20, '':MinKey} -> {'':1, '':11, '':Inf, '':MaxKey} + * {'':2, '':10, '':20, '':MinKey} -> {'':2, '':10, '':Inf, '':MaxKey} + * {'':2, '':11, '':20, '':MinKey} -> {'':2, '':11, '':Inf, '':MaxKey} + * + * TODO SERVER-48485: optimize this function to build and return the intervals as KeyString objects, + * rather than BSON. + * TODO SERVER-48473: add a query knob which sets the limit on the number of statically generated + * intervals. + */ +std::vector<std::pair<BSONObj, BSONObj>> decomposeIntoSingleIntervals( + const std::vector<OrderedIntervalList>& intervalLists, + bool lowKeyInclusive, + bool highKeyInclusive) { + invariant(intervalLists.size() > 0); + + // Appends the 'interval' bounds to the low and high keys and return the updated keys. + // Inclusivity of each bound is set through the relevant '*KeyInclusive' parameter. + auto appendInterval = [lowKeyInclusive, highKeyInclusive](const BSONObj& lowKey, + const BSONObj& highKey, + const Interval& interval) { + BSONObjBuilder lowKeyBob{lowKey}; + BSONObjBuilder highKeyBob{highKey}; + + if (interval.isMinToMax() || interval.isMaxToMin()) { + IndexBoundsBuilder::appendTrailingAllValuesInterval( + interval, lowKeyInclusive, highKeyInclusive, &lowKeyBob, &highKeyBob); + } else { + lowKeyBob.append(interval.start); + highKeyBob.append(interval.end); + } + + return std::make_pair(lowKeyBob.obj(), highKeyBob.obj()); + }; + + std::deque<std::pair<BSONObj, BSONObj>> keysQueue{{}}; + + // This is an adaptation of the BFS algorithm. The 'keysQueue' is initialized with a pair of + // empty low/high keys. For each step while traversing the 'intervalLists' we try to append the + // current interval to each generated pair in 'keysQueue' and then push the updated keys back to + // the queue. + for (auto&& list : intervalLists) { + auto size = keysQueue.size(); + for (size_t ix = 0; ix < size; ++ix) { + auto [lowKey, highKey] = keysQueue.front(); + keysQueue.pop_front(); + + for (auto&& interval : list.intervals) { + keysQueue.push_back(appendInterval(lowKey, highKey, interval)); + } + } + } + + // The 'keysQueue' contains all generated pairs of low/high keys. + return {keysQueue.begin(), keysQueue.end()}; +} + +/** + * Constructs low/high key values from the given index 'bounds if they can be represented either as + * a single interval between the low and high keys, or multiple single intervals. If index bounds + * for some interval cannot be expressed as valid low/high keys, then an empty vector is returned. + */ +std::vector<std::pair<std::unique_ptr<KeyString::Value>, std::unique_ptr<KeyString::Value>>> +makeIntervalsFromIndexBounds(const IndexBounds& bounds, + bool forward, + KeyString::Version version, + Ordering ordering) { + auto lowKeyInclusive{IndexBounds::isStartIncludedInBound(bounds.boundInclusion)}; + auto highKeyInclusive{IndexBounds::isEndIncludedInBound(bounds.boundInclusion)}; + auto intervals = [&]() -> std::vector<std::pair<BSONObj, BSONObj>> { + auto lowKey = bounds.startKey; + auto highKey = bounds.endKey; + if (bounds.isSimpleRange || + IndexBoundsBuilder::isSingleInterval( + bounds, &lowKey, &lowKeyInclusive, &highKey, &highKeyInclusive)) { + return {{lowKey, highKey}}; + } else if (canBeDecomposedIntoSingleIntervals( + bounds.fields, &lowKeyInclusive, &highKeyInclusive)) { + return decomposeIntoSingleIntervals(bounds.fields, lowKeyInclusive, highKeyInclusive); + } else { + // Index bounds cannot be represented as valid low/high keys. + return {}; + } + }(); + + LOGV2_DEBUG( + 47429005, 5, "Number of generated interval(s) for ixscan", "num"_attr = intervals.size()); + std::vector<std::pair<std::unique_ptr<KeyString::Value>, std::unique_ptr<KeyString::Value>>> + result; + for (auto&& [lowKey, highKey] : intervals) { + LOGV2_DEBUG(47429006, + 5, + "Generated interval [lowKey, highKey]", + "lowKey"_attr = lowKey, + "highKey"_attr = highKey); + // For high keys use the opposite rule as a normal seek because a forward scan should end + // after the key if inclusive, and before if exclusive. + const auto inclusive = forward != highKeyInclusive; + result.push_back({std::make_unique<KeyString::Value>( + IndexEntryComparison::makeKeyStringFromBSONKeyForSeek( + lowKey, version, ordering, forward, lowKeyInclusive)), + std::make_unique<KeyString::Value>( + IndexEntryComparison::makeKeyStringFromBSONKeyForSeek( + highKey, version, ordering, forward, inclusive))}); + } + return result; +} + +/** + * Constructs an optimized version of an index scan for multi-interval index bounds for the case + * when the bounds can be decomposed in a number of single-interval bounds. In this case, instead + * of building a generic index scan to navigate through the index using the 'IndexBoundsChecker', + * we will construct a subtree with a constant table scan containing all intervals we'd want to + * scan through. Specifically, we will build the following subtree: + * + * nlj [] [lowKeySlot, highKeySlot] + * left + * project [lowKeySlot = getField (unwindSlot, "l"), + * highKeySlot = getField (unwindSlot, "h")] + * unwind unwindSlot indexSlot boundsSlot false + * project [boundsSlot = [{"l" : KS(...), "h" : KS(...)}, + * {"l" : KS(...), "h" : KS(...)}, ...]] + * limit 1 + * coscan + * right + * ixseek lowKeySlot highKeySlot recordIdSlot [] @coll @index + * + * This subtree is similar to the single-interval subtree with the only difference that instead + * of projecting a single pair of the low/high keys, we project an array of such pairs and then + * use the unwind stage to flatten the array and generate multiple input intervals to the ixscan. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> +generateOptimizedMultiIntervalIndexScan( + const Collection* collection, + const std::string& indexName, + bool forward, + std::vector<std::pair<std::unique_ptr<KeyString::Value>, std::unique_ptr<KeyString::Value>>> + intervals, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + using namespace std::literals; + + auto recordIdSlot = slotIdGenerator->generate(); + auto lowKeySlot = slotIdGenerator->generate(); + auto highKeySlot = slotIdGenerator->generate(); + + // Construct an array containing objects with the low and high keys for each interval. E.g., + // [ {l: KS(...), h: KS(...)}, + // {l: KS(...), h: KS(...)}, ... ] + auto [boundsTag, boundsVal] = sbe::value::makeNewArray(); + auto arr = sbe::value::getArrayView(boundsVal); + for (auto&& [lowKey, highKey] : intervals) { + auto [tag, val] = sbe::value::makeNewObject(); + auto obj = sbe::value::getObjectView(val); + obj->push_back( + "l"sv, sbe::value::TypeTags::ksValue, sbe::value::bitcastFrom(lowKey.release())); + obj->push_back( + "h"sv, sbe::value::TypeTags::ksValue, sbe::value::bitcastFrom(highKey.release())); + arr->push_back(tag, val); + } + + auto boundsSlot = slotIdGenerator->generate(); + auto unwindSlot = slotIdGenerator->generate(); + + // Project out the constructed array as a constant value and add an unwind stage on top to + // flatten the array. + auto unwind = sbe::makeS<sbe::UnwindStage>( + sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + boundsSlot, + sbe::makeE<sbe::EConstant>(boundsTag, boundsVal)), + boundsSlot, + unwindSlot, + slotIdGenerator->generate(), /* We don't need an index slot but must to provide it. */ + false /* Preserve null and empty arrays, in our case it cannot be empty anyway. */); + + // Add another project stage to extract low and high keys from each value produced by unwind and + // bind the keys to the 'lowKeySlot' and 'highKeySlot'. + auto project = sbe::makeProjectStage( + std::move(unwind), + lowKeySlot, + sbe::makeE<sbe::EFunction>( + "getField"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(unwindSlot), sbe::makeE<sbe::EConstant>("l"sv))), + highKeySlot, + sbe::makeE<sbe::EFunction>("getField"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(unwindSlot), + sbe::makeE<sbe::EConstant>("h"sv)))); + + auto ixscan = sbe::makeS<sbe::IndexScanStage>( + NamespaceStringOrUUID{collection->ns().db().toString(), collection->uuid()}, + indexName, + forward, + boost::none, + recordIdSlot, + std::vector<std::string>{}, + sbe::makeSV(), + lowKeySlot, + highKeySlot, + yieldPolicy, + tracker); + + // Finally, get the keys from the outer side and feed them to the inner side (ixscan). + return {recordIdSlot, + sbe::makeS<sbe::LoopJoinStage>(std::move(project), + std::move(ixscan), + sbe::makeSV(), + sbe::makeSV(lowKeySlot, highKeySlot), + nullptr)}; +} + +/** + * Builds an anchor sub-tree of the recusrive index scan CTE to seed the result set with the initial + * 'startKey' for the index scan. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> makeAnchorBranchForGenericIndexScan( + std::unique_ptr<KeyString::Value> startKey, sbe::value::SlotIdGenerator* slotIdGenerator) { + // Just project out the 'startKey'. + auto startKeySlot = slotIdGenerator->generate(); + return {startKeySlot, + sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + startKeySlot, + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::ksValue, + sbe::value::bitcastFrom(startKey.release())))}; +} + +/** + * Builds a recursive sub-tree of the recursive CTE to generate the reminder of the result set + * consisting of valid recordId's and index seek keys to restart the index scan from. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> +makeRecursiveBranchForGenericIndexScan(const Collection* collection, + const std::string& indexName, + const sbe::CheckBoundsParams& params, + sbe::SpoolId spoolId, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + + auto resultSlot = slotIdGenerator->generate(); + auto recordIdSlot = slotIdGenerator->generate(); + auto seekKeySlot = slotIdGenerator->generate(); + auto lowKeySlot = slotIdGenerator->generate(); + + // Build a standard index scan nested loop join with the outer branch producing a low key + // to be fed into the index scan. The low key is taken from the 'seekKeySlot' which would + // contain a value from the stack spool. See below for details. + auto project = sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + lowKeySlot, + sbe::makeE<sbe::EVariable>(seekKeySlot)); + + auto ixscan = sbe::makeS<sbe::IndexScanStage>( + NamespaceStringOrUUID{collection->ns().db().toString(), collection->uuid()}, + indexName, + params.direction == 1, + resultSlot, + recordIdSlot, + std::vector<std::string>{}, + sbe::makeSV(), + lowKeySlot, + boost::none, + yieldPolicy, + tracker); + + // Get the low key from the outer side and feed it to the inner side (ixscan). + auto nlj = sbe::makeS<sbe::LoopJoinStage>( + std::move(project), std::move(ixscan), sbe::makeSV(), sbe::makeSV(lowKeySlot), nullptr); + + // Inject another nested loop join with the outer branch being a stack spool, and the inner an + // index scan nljoin which just constructed above. The stack spool is populated from the values + // generated by the index scan above, and passed through the check bounds stage, which would + // produce either a valid recordId to be consumed by the stage sitting above the index scan + // sub-tree, or a seek key to restart the index scan from. The spool will only store the seek + // keys, passing through valid recordId's. + auto checkBoundsSlot = slotIdGenerator->generate(); + return {checkBoundsSlot, + sbe::makeS<sbe::LoopJoinStage>( + sbe::makeS<sbe::SpoolConsumerStage<true>>(spoolId, sbe::makeSV(seekKeySlot)), + sbe::makeS<sbe::CheckBoundsStage>( + std::move(nlj), params, resultSlot, recordIdSlot, checkBoundsSlot), + sbe::makeSV(), + sbe::makeSV(seekKeySlot), + nullptr)}; +} + +/** + * Builds a generic multi-interval index scan for the cases when index bounds cannot be represented + * as valid low/high keys. In this case we will build a recursive sub-tree and will use the + * 'CheckBoundsStage' to navigate through the index. The recursive sub-tree is built using a union + * stage in conjunction with the stack spool: + * + * filter {isNumber(resultSlot)} + * lspool [resultSlot] {!isNumber(resultSlot)} + * union [resultSlot] + * [anchorSlot] + * project [startKeySlot = KS(...)] + * limit 1 + * coscan + * [checkBoundsSlot] + * nlj [] [seekKeySlot] + * left + * sspool [seekKeySlot] + * right + * chkbounds resultSlot recordIdSlot checkBoundsSlot + * nlj [] [lowKeySlot] + * left + * project [lowKeySlot = seekKeySlot] + * limit 1 + * coscan + * right + * ixseek lowKeySlot resultSlot recordIdSlot [] @coll @index + * + * - The anchor union branch is the starting point of the recursive subtree. It pushes the + * starting index into the lspool stage. The lspool has a filter predicate to ensure that + * only index keys will be stored in the spool. + * - There is a filter stage at the top of the subtree to ensure that we will only produce + * recordId values. + * - The recursive union branch does the remaining job. It has a nested loop join with the outer + * branch being a stack spool, which reads data from the lspool above. + * 1. The outer branch reads next seek key from sspool. + * * If the spool is empty, we're done with the scan. + * 2. The seek key is passed to the inner branch. + * 3. The inner branch execution starts with the projection of the seek key, which is + * fed into the ixscan as a 'lowKeySlot'. + * 4. Two slots produced by the ixscan, 'resultSlot' and 'recordIdSlot', are passed to + * the chkbounds stage. Note that 'resultSlot' would contain the index key. + * 5. The chkbounds stage can produce one of the following values: + * * The recordId value, taken from the ixscan stage, if the index key is within the + * bounds. + * * A seek key the ixscan will have to restart from if the key is not within the + * bounds, but has not exceeded the maximum value. + * - At this point the chkbounds stage returns ADVANCED state, to propagate the + * seek key point, but on the next call to getNext will return EOF to signal + * that we've done with the current interval and need to continue from a new + * seek point. + * * If the key is past the bound, no value is produced and EOF state is returned. + * 6. If chkbounds returns EOF, the process repeats from step 1. + * 7. Otherwise, the produces values is pulled up to the lspool stage and either is + * stored in the buffer, if it was a seek key, or just propagated to the upper stage as a + * valid recordId, and the process continues from step 4 by fetching the next key from the + * index. + * - The recursion is terminated when the sspool becomes empty. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> +generateGenericMultiIntervalIndexScan(const Collection* collection, + const IndexScanNode* ixn, + KeyString::Version version, + Ordering ordering, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::SpoolIdGenerator* spoolIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + + using namespace std::literals; + + auto resultSlot = slotIdGenerator->generate(); + + IndexBoundsChecker checker{&ixn->bounds, ixn->index.keyPattern, ixn->direction}; + IndexSeekPoint seekPoint; + + // Get the start seek key for our recursive scan. If there are no possible index entries that + // match the bounds and we cannot generate a start seek key, inject an EOF sub-tree an exit + // straight away - this index scan won't emit any results. + if (!checker.getStartSeekPoint(&seekPoint)) { + return {resultSlot, + sbe::makeS<sbe::MakeObjStage>( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 0, boost::none), + resultSlot, + boost::none, + std::vector<std::string>{}, + std::vector<std::string>{}, + sbe::makeSV(), + true, + false)}; + } + + // Build the anchor branch of the union. + auto [anchorSlot, anchorBranch] = makeAnchorBranchForGenericIndexScan( + std::make_unique<KeyString::Value>(IndexEntryComparison::makeKeyStringFromSeekPointForSeek( + seekPoint, version, ordering, ixn->direction == 1)), + slotIdGenerator); + + auto spoolId = spoolIdGenerator->generate(); + + // Build the recursive branch of the union. + auto [recursiveSlot, recursiveBranch] = makeRecursiveBranchForGenericIndexScan( + collection, + ixn->index.identifier.catalogName, + {ixn->bounds, ixn->index.keyPattern, ixn->direction, version, ordering}, + spoolId, + slotIdGenerator, + yieldPolicy, + tracker); + + // Construct a union stage from the two branches. + auto unionStage = sbe::makeS<sbe::UnionStage>( + make_vector<std::unique_ptr<sbe::PlanStage>>(std::move(anchorBranch), + std::move(recursiveBranch)), + std::vector<sbe::value::SlotVector>{sbe::makeSV(anchorSlot), sbe::makeSV(recursiveSlot)}, + sbe::makeSV(resultSlot)); + + // Stick in a lazy producer spool on top. The specified predicate will ensure that we will only + // store the seek key values in the spool (that is, if the value type is not a number, or not + // a recordId). + auto spool = sbe::makeS<sbe::SpoolLazyProducerStage>( + std::move(unionStage), + spoolId, + sbe::makeSV(resultSlot), + sbe::makeE<sbe::EPrimUnary>( + sbe::EPrimUnary::logicNot, + sbe::makeE<sbe::EFunction>("isNumber"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(resultSlot))))); + + // Finally, add a filter stage on top to filter out seek keys and return only recordIds. + return {resultSlot, + sbe::makeS<sbe::FilterStage<false>>( + std::move(spool), + sbe::makeE<sbe::EFunction>("isNumber"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(resultSlot))))}; +} +} // namespace + +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleIntervalIndexScan( + const Collection* collection, + const std::string& indexName, + bool forward, + std::unique_ptr<KeyString::Value> lowKey, + std::unique_ptr<KeyString::Value> highKey, + boost::optional<sbe::value::SlotId> recordSlot, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + auto recordIdSlot = slotIdGenerator->generate(); + auto lowKeySlot = slotIdGenerator->generate(); + auto highKeySlot = slotIdGenerator->generate(); + + // Construct a constant table scan to deliver a single row with two fields 'lowKeySlot' and + // 'highKeySlot', representing seek boundaries, into the index scan. + auto project = sbe::makeProjectStage( + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none), + lowKeySlot, + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::ksValue, + sbe::value::bitcastFrom(lowKey.release())), + highKeySlot, + sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::ksValue, + sbe::value::bitcastFrom(highKey.release()))); + + // Scan the index in the range {'lowKeySlot', 'highKeySlot'} (subject to inclusive or + // exclusive boundaries), and produce a single field recordIdSlot that can be used to + // position into the collection. + auto ixscan = sbe::makeS<sbe::IndexScanStage>( + NamespaceStringOrUUID{collection->ns().db().toString(), collection->uuid()}, + indexName, + forward, + recordSlot, + recordIdSlot, + std::vector<std::string>{}, + sbe::makeSV(), + lowKeySlot, + highKeySlot, + yieldPolicy, + tracker); + + // Finally, get the keys from the outer side and feed them to the inner side. + return {recordIdSlot, + sbe::makeS<sbe::LoopJoinStage>(std::move(project), + std::move(ixscan), + sbe::makeSV(), + sbe::makeSV(lowKeySlot, highKeySlot), + nullptr)}; +} + + +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateIndexScan( + OperationContext* opCtx, + const Collection* collection, + const IndexScanNode* ixn, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::SpoolIdGenerator* spoolIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker) { + uassert( + 4822863, "Index scans with key metadata are not supported in SBE", !ixn->addKeyMetadata); + uassert(4822864, "Index scans with a filter are not supported in SBE", !ixn->filter); + + auto descriptor = + collection->getIndexCatalog()->findIndexByName(opCtx, ixn->index.identifier.catalogName); + auto accessMethod = collection->getIndexCatalog()->getEntry(descriptor)->accessMethod(); + auto intervals = + makeIntervalsFromIndexBounds(ixn->bounds, + ixn->direction == 1, + accessMethod->getSortedDataInterface()->getKeyStringVersion(), + accessMethod->getSortedDataInterface()->getOrdering()); + + auto [slot, stage] = [&]() { + if (intervals.size() == 1) { + // If we have just a single interval, we can construct a simplified sub-tree. + auto&& [lowKey, highKey] = intervals[0]; + return generateSingleIntervalIndexScan(collection, + ixn->index.identifier.catalogName, + ixn->direction == 1, + std::move(lowKey), + std::move(highKey), + boost::none, + slotIdGenerator, + yieldPolicy, + tracker); + } else if (intervals.size() > 1) { + // Or, if we were able to decompose multi-interval index bounds into a number of + // single-interval bounds, we can also built an optimized sub-tree to perform an index + // scan. + return generateOptimizedMultiIntervalIndexScan(collection, + ixn->index.identifier.catalogName, + ixn->direction == 1, + std::move(intervals), + slotIdGenerator, + yieldPolicy, + tracker); + } else { + // Otherwise, build a generic index scan for multi-interval index bounds. + return generateGenericMultiIntervalIndexScan( + collection, + ixn, + accessMethod->getSortedDataInterface()->getKeyStringVersion(), + accessMethod->getSortedDataInterface()->getOrdering(), + slotIdGenerator, + spoolIdGenerator, + yieldPolicy, + tracker); + } + }(); + + if (ixn->shouldDedup) { + stage = sbe::makeS<sbe::HashAggStage>(std::move(stage), sbe::makeSV(slot), sbe::makeEM()); + } + + return {slot, std::move(stage)}; +} +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.h b/src/mongo/db/query/sbe_stage_builder_index_scan.h new file mode 100644 index 00000000000..acc9e8be7b4 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_index_scan.h @@ -0,0 +1,77 @@ +/** + * 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/exec/trial_run_progress_tracker.h" +#include "mongo/db/query/query_solution.h" + +namespace mongo::stage_builder { +/** + * Generates an SBE plan stage sub-tree implementing an index scan. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateIndexScan( + OperationContext* opCtx, + const Collection* collection, + const IndexScanNode* ixn, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::SpoolIdGenerator* spoolIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker); + +/** + * Constructs the most simple version of an index scan from the single interval index bounds. The + * generated subtree will have the following form: + * + * nlj [] [lowKeySlot, highKeySlot] + * left + * project [lowKeySlot = KS(...), highKeySlot = KS(...)] + * limit 1 + * coscan + * right + * ixseek lowKeySlot highKeySlot recordIdSlot [] @coll @index + * + * The inner branch of the nested loop join produces a single row with the low/high keys which is + * fed to the ixscan. + * + * If 'recordSlot' is provided, than the corresponding slot will be filled out with each KeyString + * in the index. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleIntervalIndexScan( + const Collection* collection, + const std::string& indexName, + bool forward, + std::unique_ptr<KeyString::Value> lowKey, + std::unique_ptr<KeyString::Value> highKey, + boost::optional<sbe::value::SlotId> recordSlot, + sbe::value::SlotIdGenerator* slotIdGenerator, + PlanYieldPolicy* yieldPolicy, + TrialRunProgressTracker* tracker); +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_projection.cpp b/src/mongo/db/query/sbe_stage_builder_projection.cpp new file mode 100644 index 00000000000..39ca0bd89e2 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_projection.cpp @@ -0,0 +1,382 @@ +/** + * 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/query/sbe_stage_builder_projection.h" + +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/stages/makeobj.h" +#include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/stages/traverse.h" +#include "mongo/db/exec/sbe/values/bson.h" +#include "mongo/db/query/sbe_stage_builder_expression.h" +#include "mongo/db/query/tree_walker.h" +#include "mongo/util/str.h" +#include "mongo/util/visit_helper.h" + +namespace mongo::stage_builder { +namespace { +using ExpressionType = std::unique_ptr<sbe::EExpression>; +using PlanStageType = std::unique_ptr<sbe::PlanStage>; + +/** + * Stores context across calls to visit() in the projection traversal visitors. + */ +struct ProjectionTraversalVisitorContext { + // Stores the field names to visit at each nested projection level, and the base path to this + // level. Top of the stack is the most recently visited level. + struct NestedLevel { + // The input slot for the current level. This is the parent sub-document for each of the + // projected fields at the current level. + sbe::value::SlotId inputSlot; + // The fields names at the current projection level. + std::list<std::string> fields; + // All but the last path component of the current path being visited. None if at the + // top-level and there is no "parent" path. + std::stack<std::string> basePath; + // A traversal sub-tree which combines traversals for each of the fields at the current + // level. + PlanStageType fieldPathExpressionsTraverseStage{ + sbe::makeS<sbe::LimitSkipStage>(sbe::makeS<sbe::CoScanStage>(), 1, boost::none)}; + }; + + // Stores evaluation expressions for each of the projections at the current nested level. It can + // evaluate either to an expression, if we're at the leaf node of the projection, or a sub-tree, + // if we're evaluating a field projection in the middle of the path. The stack elements are + // optional because for an exclusion projection we don't need to evaluate anything, but we need + // to have an element on the stack which corresponds to a projected field. + struct ProjectEval { + sbe::value::SlotId inputSlot; + sbe::value::SlotId outputSlot; + stdx::variant<PlanStageType, ExpressionType> expr; + }; + + const auto& topFrontField() const { + invariant(!levels.empty()); + invariant(!levels.top().fields.empty()); + return levels.top().fields.front(); + } + + void popFrontField() { + invariant(!levels.empty()); + invariant(!levels.top().fields.empty()); + levels.top().fields.pop_front(); + } + + auto& topLevel() { + invariant(!levels.empty()); + return levels.top(); + } + + void popLevel() { + invariant(!levels.empty()); + invariant(levels.top().fields.empty()); + levels.pop(); + } + + void pushLevel(std::list<std::string> fields) { + levels.push({levels.empty() ? inputSlot : slotIdGenerator->generate(), std::move(fields)}); + } + + std::pair<sbe::value::SlotId, PlanStageType> done() { + invariant(evals.size() == 1); + auto eval = std::move(evals.top()); + invariant(eval); + invariant(stdx::holds_alternative<PlanStageType>(eval->expr)); + return {eval->outputSlot, + sbe::makeS<sbe::TraverseStage>(std::move(inputStage), + std::move(stdx::get<PlanStageType>(eval->expr)), + eval->inputSlot, + eval->outputSlot, + eval->outputSlot, + sbe::makeSV(), + nullptr, + nullptr)}; + } + + projection_ast::ProjectType projectType; + sbe::value::SlotIdGenerator* const slotIdGenerator; + sbe::value::FrameIdGenerator* const frameIdGenerator; + + // The input stage to this projection and the slot to read a root document from. + PlanStageType inputStage; + sbe::value::SlotId inputSlot; + std::stack<NestedLevel> levels; + std::stack<boost::optional<ProjectEval>> evals; + + // See the comment above the generateExpression() declaration for an explanation of the + // 'relevantSlots' list. + sbe::value::SlotVector relevantSlots; +}; + +/** + * A projection traversal pre-visitor used for maintaining nested levels while traversing a + * projection AST. + */ +class ProjectionTraversalPreVisitor final : public projection_ast::ProjectionASTConstVisitor { +public: + ProjectionTraversalPreVisitor(ProjectionTraversalVisitorContext* context) : _context{context} { + invariant(_context); + } + + void visit(const projection_ast::ProjectionPathASTNode* node) final { + if (node->parent()) { + _context->topLevel().basePath.push(_context->topFrontField()); + _context->popFrontField(); + } + _context->pushLevel({node->fieldNames().begin(), node->fieldNames().end()}); + } + + void visit(const projection_ast::ProjectionPositionalASTNode* node) final { + uasserted(4822885, str::stream() << "Positional projection is not supported in SBE"); + } + + void visit(const projection_ast::ProjectionSliceASTNode* node) final { + uasserted(4822886, str::stream() << "Slice projection is not supported in SBE"); + } + + void visit(const projection_ast::ProjectionElemMatchASTNode* node) final { + uasserted(4822887, str::stream() << "ElemMatch projection is not supported in SBE"); + } + + void visit(const projection_ast::ExpressionASTNode* node) final {} + + void visit(const projection_ast::MatchExpressionASTNode* node) final { + uasserted(4822888, + str::stream() << "Projection match expressions are not supported in SBE"); + } + + void visit(const projection_ast::BooleanConstantASTNode* node) final {} + +private: + ProjectionTraversalVisitorContext* _context; +}; + +/** + * A projection traversal post-visitor used for maintaining nested levels while traversing a + * projection AST and producing an SBE traversal sub-tree for each nested level. + */ +class ProjectionTraversalPostVisitor final : public projection_ast::ProjectionASTConstVisitor { +public: + ProjectionTraversalPostVisitor(ProjectionTraversalVisitorContext* context) : _context{context} { + invariant(_context); + } + + void visit(const projection_ast::BooleanConstantASTNode* node) final { + using namespace std::literals; + + // If this is an inclusion projection, extract the field and push a getField expression on + // top of the 'evals' stack. For an exclusion projection just push an empty optional. + if (node->value()) { + _context->evals.push( + {{_context->topLevel().inputSlot, + _context->slotIdGenerator->generate(), + sbe::makeE<sbe::EFunction>( + "getField"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>(_context->topLevel().inputSlot), + sbe::makeE<sbe::EConstant>(_context->topFrontField())))}}); + } else { + _context->evals.push({}); + } + _context->popFrontField(); + } + + void visit(const projection_ast::ExpressionASTNode* node) final { + // Generate an expression to evaluate a projection expression and push it on top of the + // 'evals' stack. If the expression is translated into a sub-tree, stack it with the + // existing 'fieldPathExpressionsTraverseStage' sub-tree. + auto [outputSlot, expr, stage] = + generateExpression(node->expressionRaw(), + std::move(_context->topLevel().fieldPathExpressionsTraverseStage), + _context->slotIdGenerator, + _context->frameIdGenerator, + _context->inputSlot, + &_context->relevantSlots); + _context->evals.push({{_context->topLevel().inputSlot, outputSlot, std::move(expr)}}); + _context->topLevel().fieldPathExpressionsTraverseStage = std::move(stage); + _context->popFrontField(); + } + + void visit(const projection_ast::ProjectionPathASTNode* node) final { + using namespace std::literals; + + const auto isInclusion = _context->projectType == projection_ast::ProjectType::kInclusion; + sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projects; + sbe::value::SlotVector projectSlots; + std::vector<std::string> projectFields; + std::vector<std::string> restrictFields; + auto inputStage{std::move(_context->topLevel().fieldPathExpressionsTraverseStage)}; + + invariant(_context->evals.size() >= node->fieldNames().size()); + // Walk through all the fields at the current nested level in reverse order (to match field + // names with the elements on the 'evals' stack) and, + // * For exclusion projections populate the 'restrictFields' array to be passed to the + // mkobj stage below, which constructs an output document for the current nested level. + // * For inclusion projections, + // - Populates 'projectFields' and 'projectSlots' vectors holding field names to + // project, and slots to access evaluated projection values. + // - Populates 'projects' map to actually project out the values. + // - For nested paths injects a traversal sub-tree. + for (auto it = node->fieldNames().rbegin(); it != node->fieldNames().rend(); ++it) { + auto eval = std::move(_context->evals.top()); + _context->evals.pop(); + + // If the projection eval element is empty, then this is an exclusion projection and we + // can put the field name to the vector of restricted fields. + if (!eval) { + restrictFields.push_back(*it); + continue; + } + + projectSlots.push_back(eval->outputSlot); + projectFields.push_back(*it); + + stdx::visit( + visit_helper::Overloaded{ + [&](ExpressionType& expr) { + projects.emplace(eval->outputSlot, std::move(expr)); + }, + [&](PlanStageType& stage) { + invariant(!_context->topLevel().basePath.empty()); + + inputStage = sbe::makeProjectStage( + std::move(inputStage), + eval->inputSlot, + sbe::makeE<sbe::EFunction>( + "getField"sv, + sbe::makeEs( + sbe::makeE<sbe::EVariable>(_context->topLevel().inputSlot), + sbe::makeE<sbe::EConstant>( + _context->topLevel().basePath.top())))); + _context->topLevel().basePath.pop(); + + inputStage = sbe::makeS<sbe::TraverseStage>(std::move(inputStage), + std::move(stage), + eval->inputSlot, + eval->outputSlot, + eval->outputSlot, + sbe::makeSV(), + nullptr, + nullptr); + }}, + eval->expr); + } + + // We walked through the field names in reverse order, so need to reverse the following two + // vectors. + std::reverse(projectFields.begin(), projectFields.end()); + std::reverse(projectSlots.begin(), projectSlots.end()); + + // If we have something to actually project, then inject a projection stage. + if (!projects.empty()) { + inputStage = sbe::makeS<sbe::ProjectStage>(std::move(inputStage), std::move(projects)); + } + + // Finally, inject an mkobj stage to generate a document for the current nested level. For + // inclusion projection also add constant filter stage on top to filter out input values for + // nested traversal if they're not documents. + auto outputSlot = _context->slotIdGenerator->generate(); + _context->relevantSlots.push_back(outputSlot); + _context->evals.push( + {{_context->topLevel().inputSlot, + outputSlot, + isInclusion ? sbe::makeS<sbe::FilterStage<true>>( + sbe::makeS<sbe::MakeObjStage>(std::move(inputStage), + outputSlot, + boost::none, + std::vector<std::string>{}, + std::move(projectFields), + std::move(projectSlots), + true, + false), + sbe::makeE<sbe::EFunction>("isObject"sv, + sbe::makeEs(sbe::makeE<sbe::EVariable>( + _context->topLevel().inputSlot)))) + : sbe::makeS<sbe::MakeObjStage>(std::move(inputStage), + outputSlot, + _context->topLevel().inputSlot, + std::move(restrictFields), + std::move(projectFields), + std::move(projectSlots), + false, + true)}}); + // We've done with the current nested level. + _context->popLevel(); + } + + void visit(const projection_ast::ProjectionPositionalASTNode* node) final {} + void visit(const projection_ast::ProjectionSliceASTNode* node) final {} + void visit(const projection_ast::ProjectionElemMatchASTNode* node) final {} + void visit(const projection_ast::MatchExpressionASTNode* node) final {} + +private: + ProjectionTraversalVisitorContext* _context; +}; + +class ProjectionTraversalWalker final { +public: + ProjectionTraversalWalker(projection_ast::ProjectionASTConstVisitor* preVisitor, + projection_ast::ProjectionASTConstVisitor* postVisitor) + : _preVisitor{preVisitor}, _postVisitor{postVisitor} {} + + void preVisit(const projection_ast::ASTNode* node) { + node->acceptVisitor(_preVisitor); + } + + void postVisit(const projection_ast::ASTNode* node) { + node->acceptVisitor(_postVisitor); + } + + void inVisit(long count, const projection_ast::ASTNode* node) {} + +private: + projection_ast::ProjectionASTConstVisitor* _preVisitor; + projection_ast::ProjectionASTConstVisitor* _postVisitor; +}; +} // namespace + +std::pair<sbe::value::SlotId, PlanStageType> generateProjection( + const projection_ast::Projection* projection, + PlanStageType stage, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::FrameIdGenerator* frameIdGenerator, + sbe::value::SlotId inputVar) { + ProjectionTraversalVisitorContext context{ + projection->type(), slotIdGenerator, frameIdGenerator, std::move(stage), inputVar}; + context.relevantSlots.push_back(inputVar); + ProjectionTraversalPreVisitor preVisitor{&context}; + ProjectionTraversalPostVisitor postVisitor{&context}; + ProjectionTraversalWalker walker{&preVisitor, &postVisitor}; + tree_walker::walk<true, projection_ast::ASTNode>(projection->root(), &walker); + return context.done(); +} +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_projection.h b/src/mongo/db/query/sbe_stage_builder_projection.h new file mode 100644 index 00000000000..5ac24b10683 --- /dev/null +++ b/src/mongo/db/query/sbe_stage_builder_projection.h @@ -0,0 +1,48 @@ +/** + * 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/projection.h" + +namespace mongo::stage_builder { +/** + * Generates an SBE plan stage sub-tree implementing a query projection. The 'stage' parameter + * defines an input stage to the generated SBE plan stage sub-tree. The 'inputVar' defines a + * variable to read the input document from. + */ +std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateProjection( + const projection_ast::Projection* proj, + std::unique_ptr<sbe::PlanStage> stage, + sbe::value::SlotIdGenerator* slotIdGenerator, + sbe::value::FrameIdGenerator* frameIdGenerator, + sbe::value::SlotId inputVar); + +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_sub_planner.cpp b/src/mongo/db/query/sbe_sub_planner.cpp new file mode 100644 index 00000000000..4ecdc93344a --- /dev/null +++ b/src/mongo/db/query/sbe_sub_planner.cpp @@ -0,0 +1,114 @@ +/** + * 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/query/sbe_sub_planner.h" + +#include "mongo/db/query/collection_query_info.h" +#include "mongo/db/query/query_planner.h" +#include "mongo/db/query/sbe_multi_planner.h" +#include "mongo/db/query/stage_builder_util.h" + +namespace mongo::sbe { +plan_ranker::CandidatePlan SubPlanner::plan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) { + // Plan each branch of the $or. + auto subplanningStatus = + QueryPlanner::planSubqueries(_opCtx, + _collection, + CollectionQueryInfo::get(_collection).getPlanCache(), + _cq, + _queryParams); + if (!subplanningStatus.isOK()) { + return planWholeQuery(); + } + + auto multiplanCallback = [&](CanonicalQuery* cq, + std::vector<std::unique_ptr<QuerySolution>> solutions) + -> StatusWith<std::unique_ptr<QuerySolution>> { + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots; + for (auto&& solution : solutions) { + roots.push_back(stage_builder::buildSlotBasedExecutableTree( + _opCtx, _collection, *cq, *solution, _yieldPolicy, true)); + } + + // We pass the SometimesCache option to the MPS because the SubplanStage currently does + // not use the 'CachedSolutionPlanner' eviction mechanism. We therefore are more + // conservative about putting a potentially bad plan into the cache in the subplan path. + MultiPlanner multiPlanner{ + _opCtx, _collection, *cq, PlanCachingMode::SometimesCache, _yieldPolicy}; + auto plan = multiPlanner.plan(std::move(solutions), std::move(roots)); + return std::move(plan.solution); + }; + + auto subplanSelectStat = QueryPlanner::choosePlanForSubqueries( + _cq, _queryParams, std::move(subplanningStatus.getValue()), multiplanCallback); + if (!subplanSelectStat.isOK()) { + // 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 (an index + // may have been dropped, we may have exceeded the time limit, etc). + uassert(4822874, + subplanSelectStat.getStatus().reason(), + subplanSelectStat == ErrorCodes::NoQueryExecutionPlans); + return planWholeQuery(); + } + + // Build a plan stage tree from a composite solution. + auto compositeSolution = std::move(subplanSelectStat.getValue()); + auto&& [root, data] = stage_builder::buildSlotBasedExecutableTree( + _opCtx, _collection, _cq, *compositeSolution, _yieldPolicy, false); + prepareExecutionPlan(root.get(), &data); + return {std::move(compositeSolution), std::move(root), std::move(data)}; +} + +plan_ranker::CandidatePlan SubPlanner::planWholeQuery() const { + // Use the query planning module to plan the whole query. + auto solutions = uassertStatusOK(QueryPlanner::plan(_cq, _queryParams)); + + // Only one possible plan. Build the stages from the solution. + if (solutions.size() == 1) { + auto&& [root, data] = stage_builder::buildSlotBasedExecutableTree( + _opCtx, _collection, _cq, *solutions[0], _yieldPolicy, false); + prepareExecutionPlan(root.get(), &data); + return {std::move(solutions[0]), std::move(root), std::move(data)}; + } + + // Many solutions. Build a plan stage tree for each solution and create a multi planner to pick + // the best, update the cache, and so on. + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots; + for (auto&& solution : solutions) { + roots.push_back(stage_builder::buildSlotBasedExecutableTree( + _opCtx, _collection, _cq, *solution, _yieldPolicy, true)); + } + + MultiPlanner multiPlanner{_opCtx, _collection, _cq, PlanCachingMode::AlwaysCache, _yieldPolicy}; + return multiPlanner.plan(std::move(solutions), std::move(roots)); +} +} // namespace mongo::sbe diff --git a/src/mongo/db/query/sbe_sub_planner.h b/src/mongo/db/query/sbe_sub_planner.h new file mode 100644 index 00000000000..ef6e6ae3e35 --- /dev/null +++ b/src/mongo/db/query/sbe_sub_planner.h @@ -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. + */ + +#pragma once + +#include "mongo/db/query/sbe_plan_ranker.h" +#include "mongo/db/query/sbe_runtime_planner.h" + +namespace mongo::sbe { +/** + * This runtime planner is used for rooted $or queries. It plans each clause of the $or + * individually, and then creates an overall query plan based on the winning plan from each + * clause. + * + * Uses the 'MultiPlanner' in order to pick best plans for the individual clauses. + */ +class SubPlanner final : public BaseRuntimePlanner { +public: + SubPlanner(OperationContext* opCtx, + Collection* collection, + const CanonicalQuery& cq, + const QueryPlannerParams& queryParams, + PlanYieldPolicySBE* yieldPolicy) + : BaseRuntimePlanner{opCtx, collection, cq, yieldPolicy}, _queryParams{queryParams} {} + + plan_ranker::CandidatePlan plan( + std::vector<std::unique_ptr<QuerySolution>> solutions, + std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots) + final; + +private: + plan_ranker::CandidatePlan planWholeQuery() const; + + // Query parameters used to create a query solution for each $or branch. + const QueryPlannerParams _queryParams; +}; +} // namespace mongo::sbe diff --git a/src/mongo/db/query/stage_builder.h b/src/mongo/db/query/stage_builder.h index a0437832822..dbee0b6b8b6 100644 --- a/src/mongo/db/query/stage_builder.h +++ b/src/mongo/db/query/stage_builder.h @@ -29,31 +29,35 @@ #pragma once -#include "mongo/db/exec/plan_stage.h" #include "mongo/db/exec/working_set.h" #include "mongo/db/query/query_solution.h" -namespace mongo { - -class OperationContext; - +namespace mongo::stage_builder { /** - * The StageBuilder converts a QuerySolution to an executable tree of PlanStage(s). + * The StageBuilder converts a QuerySolution tree to an executable tree of PlanStage(s), with the + * specific type defined by the 'PlanStageType' parameter. */ +template <typename PlanStageType> class StageBuilder { public: + StageBuilder(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + const QuerySolution& solution) + : _opCtx(opCtx), _collection(collection), _cq(cq), _solution(solution) {} + + virtual ~StageBuilder() = default; + /** - * Turns 'solution' into an executable tree of PlanStage(s). Returns a pointer to the root of - * the plan stage tree. - * - * 'cq' must be the CanonicalQuery from which 'solution' is derived. Illegal to call if 'wsIn' - * is nullptr, or if 'solution.root' is nullptr. + * Given a root node of a QuerySolution tree, builds and returns a corresponding executable + * tree of PlanStages. */ - static std::unique_ptr<PlanStage> build(OperationContext* opCtx, - const Collection* collection, - const CanonicalQuery& cq, - const QuerySolution& solution, - WorkingSet* wsIn); -}; + virtual std::unique_ptr<PlanStageType> build(const QuerySolutionNode* root) = 0; -} // namespace mongo +protected: + OperationContext* _opCtx; + const Collection* _collection; + const CanonicalQuery& _cq; + const QuerySolution& _solution; +}; +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/stage_builder_util.cpp b/src/mongo/db/query/stage_builder_util.cpp new file mode 100644 index 00000000000..5f7c6c462e3 --- /dev/null +++ b/src/mongo/db/query/stage_builder_util.cpp @@ -0,0 +1,78 @@ +/** + * 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/query/stage_builder_util.h" + +#include "mongo/db/query/classic_stage_builder.h" +#include "mongo/db/query/plan_yield_policy.h" +#include "mongo/db/query/sbe_stage_builder.h" + +namespace mongo::stage_builder { +std::unique_ptr<PlanStage> buildClassicExecutableTree(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + const QuerySolution& solution, + WorkingSet* ws) { + // Only QuerySolutions derived from queries parsed with context, or QuerySolutions derived from + // queries that disallow extensions, can be properly executed. If the query does not have + // $text/$where context (and $text/$where are allowed), then no attempt should be made to + // execute the query. + invariant(!cq.canHaveNoopMatchNodes()); + invariant(solution.root); + invariant(ws); + auto builder = std::make_unique<ClassicStageBuilder>(opCtx, collection, cq, solution, ws); + return builder->build(solution.root.get()); +} + +std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> +buildSlotBasedExecutableTree(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + const QuerySolution& solution, + PlanYieldPolicy* yieldPolicy, + bool needsTrialRunProgressTracker) { + // Only QuerySolutions derived from queries parsed with context, or QuerySolutions derived from + // queries that disallow extensions, can be properly executed. If the query does not have + // $text/$where context (and $text/$where are allowed), then no attempt should be made to + // execute the query. + invariant(!cq.canHaveNoopMatchNodes()); + invariant(solution.root); + + auto sbeYieldPolicy = dynamic_cast<PlanYieldPolicySBE*>(yieldPolicy); + invariant(sbeYieldPolicy); + + auto builder = std::make_unique<SlotBasedStageBuilder>( + opCtx, collection, cq, solution, sbeYieldPolicy, needsTrialRunProgressTracker); + auto root = builder->build(solution.root.get()); + auto data = builder->getPlanStageData(); + return {std::move(root), std::move(data)}; +} +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/stage_builder_util.h b/src/mongo/db/query/stage_builder_util.h new file mode 100644 index 00000000000..aa6c0abfad2 --- /dev/null +++ b/src/mongo/db/query/stage_builder_util.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/query/classic_stage_builder.h" +#include "mongo/db/query/plan_yield_policy.h" +#include "mongo/db/query/sbe_stage_builder.h" + +namespace mongo::stage_builder { +/** + * Turns 'solution' into an executable tree of PlanStage(s). Returns a pointer to the root of + * the plan stage tree. + * + * 'cq' must be the CanonicalQuery from which 'solution' is derived. Illegal to call if 'ws' + * is nullptr, or if 'solution.root' is nullptr. + * + * The 'PlanStageType' type parameter defines a specific type of PlanStage the executable tree + * will consist of. + */ +std::unique_ptr<PlanStage> buildClassicExecutableTree(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + const QuerySolution& solution, + WorkingSet* ws); + +std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> +buildSlotBasedExecutableTree(OperationContext* opCtx, + const Collection* collection, + const CanonicalQuery& cq, + const QuerySolution& solution, + PlanYieldPolicy* yieldPolicy, + bool needsTrialRunProgressTracker); + +} // namespace mongo::stage_builder diff --git a/src/mongo/db/query/projection_ast_walker.h b/src/mongo/db/query/tree_walker.h index 64674dbd3d0..eb36a27e6ca 100644 --- a/src/mongo/db/query/projection_ast_walker.h +++ b/src/mongo/db/query/tree_walker.h @@ -29,40 +29,45 @@ #pragma once -#include <boost/intrusive_ptr.hpp> - -#include "mongo/db/query/projection_ast.h" - -namespace mongo::projection_ast_walker { - +namespace mongo::tree_walker { +/** + * A template type which resolves to 'const T*' if 'IsConst' argument is 'true', and to 'T*' + * otherwise. + */ +template <bool IsConst, typename T> +using MaybeConstPtr = typename std::conditional<IsConst, const T*, T*>::type; /** - * Provided with a Walker and an ASTNode*, walk() calls each of the following: + * Provided with a Walker and a Node, walk() calls each of the following: * * walker.preVisit() once before walking to each child. * * walker.inVisit() between walking to each child. It is called multiple times, once between each - * pair of children. walker.inVisit() is skipped if the ASTNode has fewer than two children. + * pair of children. walker.inVisit() is skipped if the Node has fewer than two children. * * walker.postVisit() once after walking to each child. - * Each of the ASTNode's child ASTNode is recursively walked and the same three methods are - * called for it. * - * If the caller doesn't intend to modify the AST, then the template argument 'IsConst' should be + * Each of the Node's children is recursively walked and the same three methods are called for it. + * + * The Node type should either provide begin() and end() methods returning an iterator to walk its + * children, or define begin() and end() functions taking a Node reference which return the + * iterator. + * + * If the caller doesn't intend to modify the tree, then the template argument 'IsConst' should be * set to 'true'. In this case the 'node' pointer will be qualified with 'const'. */ -template <typename Walker, bool IsConst = true> -void walk(Walker* walker, projection_ast::MaybeConstPtr<IsConst, projection_ast::ASTNode> node) { +template <bool IsConst, typename Node, typename Walker> +void walk(MaybeConstPtr<IsConst, Node> node, Walker* walker) { if (node) { walker->preVisit(node); auto count = 0ull; - for (auto&& child : node->children()) { - if (count) + for (auto&& child : *node) { + if (count) { walker->inVisit(count, node); + } ++count; - walk<Walker, IsConst>(walker, child.get()); + walk<IsConst, Node, Walker>(&*child, walker); } walker->postVisit(node); } } - -} // namespace mongo::projection_ast_walker +} // namespace mongo::tree_walker diff --git a/src/mongo/db/repl/dbcheck.cpp b/src/mongo/db/repl/dbcheck.cpp index 083dd74edc1..c2b30155333 100644 --- a/src/mongo/db/repl/dbcheck.cpp +++ b/src/mongo/db/repl/dbcheck.cpp @@ -194,7 +194,7 @@ DbCheckHasher::DbCheckHasher(OperationContext* opCtx, start.obj(), end.obj(), BoundInclusion::kIncludeEndKeyOnly, - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, InternalPlanner::FORWARD, InternalPlanner::IXSCAN_FETCH); } diff --git a/src/mongo/db/repl/idempotency_test_fixture.cpp b/src/mongo/db/repl/idempotency_test_fixture.cpp index fc54c2f3933..07aa9d39d3c 100644 --- a/src/mongo/db/repl/idempotency_test_fixture.cpp +++ b/src/mongo/db/repl/idempotency_test_fixture.cpp @@ -333,7 +333,7 @@ std::string IdempotencyTest::computeDataHash(Collection* collection) { BSONObj(), BSONObj(), BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, InternalPlanner::FORWARD, InternalPlanner::IXSCAN_FETCH); ASSERT(nullptr != exec.get()); diff --git a/src/mongo/db/repl/oplog_applier_impl_test_fixture.cpp b/src/mongo/db/repl/oplog_applier_impl_test_fixture.cpp index 45219111f34..07fd4a30ee7 100644 --- a/src/mongo/db/repl/oplog_applier_impl_test_fixture.cpp +++ b/src/mongo/db/repl/oplog_applier_impl_test_fixture.cpp @@ -323,7 +323,7 @@ CollectionReader::CollectionReader(OperationContext* opCtx, const NamespaceStrin _exec(InternalPlanner::collectionScan(opCtx, nss.ns(), _collToScan.getCollection(), - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, InternalPlanner::FORWARD)) {} StatusWith<BSONObj> CollectionReader::next() { diff --git a/src/mongo/db/repl/oplog_interface_local.cpp b/src/mongo/db/repl/oplog_interface_local.cpp index d2bbdbb5f4e..a0b9e89a404 100644 --- a/src/mongo/db/repl/oplog_interface_local.cpp +++ b/src/mongo/db/repl/oplog_interface_local.cpp @@ -63,7 +63,7 @@ OplogIteratorLocal::OplogIteratorLocal(OperationContext* opCtx) NamespaceString::kRsOplogNamespace.ns(), CollectionCatalog::get(opCtx).lookupCollectionByNamespace( opCtx, NamespaceString::kRsOplogNamespace), - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, InternalPlanner::BACKWARD)) {} StatusWith<OplogInterface::Iterator::Value> OplogIteratorLocal::next() { diff --git a/src/mongo/db/repl/replication_info.cpp b/src/mongo/db/repl/replication_info.cpp index d9f0d7d6602..94075c12c63 100644 --- a/src/mongo/db/repl/replication_info.cpp +++ b/src/mongo/db/repl/replication_info.cpp @@ -145,8 +145,10 @@ TopologyVersion appendReplicationInfo(OperationContext* opCtx, { const NamespaceString localSources{"local.sources"}; AutoGetCollectionForReadCommand ctx(opCtx, localSources); - auto exec = InternalPlanner::collectionScan( - opCtx, localSources.ns(), ctx.getCollection(), PlanExecutor::NO_YIELD); + auto exec = InternalPlanner::collectionScan(opCtx, + localSources.ns(), + ctx.getCollection(), + PlanYieldPolicy::YieldPolicy::NO_YIELD); BSONObj obj; PlanExecutor::ExecState state; while (PlanExecutor::ADVANCED == (state = exec->getNext(&obj, nullptr))) { diff --git a/src/mongo/db/repl/rollback_impl.cpp b/src/mongo/db/repl/rollback_impl.cpp index 9c13c0e0fe2..83ab9d39004 100644 --- a/src/mongo/db/repl/rollback_impl.cpp +++ b/src/mongo/db/repl/rollback_impl.cpp @@ -608,8 +608,9 @@ void RollbackImpl::_correctRecordStoreCounts(OperationContext* opCtx) { invariant(coll == collToScan, str::stream() << "Catalog returned invalid collection: " << nss.ns() << " (" << uuid.toString() << ")"); - auto exec = collToScan->makePlanExecutor( - opCtx, PlanExecutor::INTERRUPT_ONLY, Collection::ScanDirection::kForward); + auto exec = collToScan->makePlanExecutor(opCtx, + PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY, + Collection::ScanDirection::kForward); long long countFromScan = 0; PlanExecutor::ExecState state; while (PlanExecutor::ADVANCED == diff --git a/src/mongo/db/repl/rs_rollback.cpp b/src/mongo/db/repl/rs_rollback.cpp index 20c2d06f699..5aa58b0ee96 100644 --- a/src/mongo/db/repl/rs_rollback.cpp +++ b/src/mongo/db/repl/rs_rollback.cpp @@ -1030,8 +1030,7 @@ void dropCollection(OperationContext* opCtx, // Performs a collection scan and writes all documents in the collection to disk // in order to keep an archive of items that were rolled back. auto exec = InternalPlanner::collectionScan( - opCtx, nss.toString(), collection, PlanExecutor::YIELD_AUTO); - + opCtx, nss.toString(), collection, PlanYieldPolicy::YieldPolicy::YIELD_AUTO); PlanExecutor::ExecState execState; try { BSONObj curObj; diff --git a/src/mongo/db/repl/storage_interface_impl.cpp b/src/mongo/db/repl/storage_interface_impl.cpp index ec9bdd75924..61f9547cf7f 100644 --- a/src/mongo/db/repl/storage_interface_impl.cpp +++ b/src/mongo/db/repl/storage_interface_impl.cpp @@ -658,13 +658,16 @@ StatusWith<std::vector<BSONObj>> _findOrDeleteDocuments( } // Use collection scan. planExecutor = isFind - ? InternalPlanner::collectionScan( - opCtx, nsOrUUID.toString(), collection, PlanExecutor::NO_YIELD, direction) + ? InternalPlanner::collectionScan(opCtx, + nsOrUUID.toString(), + collection, + PlanYieldPolicy::YieldPolicy::NO_YIELD, + direction) : InternalPlanner::deleteWithCollectionScan( opCtx, collection, makeDeleteStageParamsForDeleteDocuments(), - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, direction); } else { // Use index scan. @@ -696,25 +699,26 @@ StatusWith<std::vector<BSONObj>> _findOrDeleteDocuments( if (!endKey.isEmpty()) { bounds.second = endKey; } - planExecutor = isFind ? InternalPlanner::indexScan(opCtx, - collection, - indexDescriptor, - bounds.first, - bounds.second, - boundInclusion, - PlanExecutor::NO_YIELD, - direction, - InternalPlanner::IXSCAN_FETCH) - : InternalPlanner::deleteWithIndexScan( - opCtx, - collection, - makeDeleteStageParamsForDeleteDocuments(), - indexDescriptor, - bounds.first, - bounds.second, - boundInclusion, - PlanExecutor::NO_YIELD, - direction); + planExecutor = isFind + ? InternalPlanner::indexScan(opCtx, + collection, + indexDescriptor, + bounds.first, + bounds.second, + boundInclusion, + PlanYieldPolicy::YieldPolicy::NO_YIELD, + direction, + InternalPlanner::IXSCAN_FETCH) + : InternalPlanner::deleteWithIndexScan( + opCtx, + collection, + makeDeleteStageParamsForDeleteDocuments(), + indexDescriptor, + bounds.first, + bounds.second, + boundInclusion, + PlanYieldPolicy::YieldPolicy::NO_YIELD, + direction); } std::vector<BSONObj> docs; @@ -867,7 +871,7 @@ Status _updateWithQuery(OperationContext* opCtx, const Timestamp& ts) { invariant(!request.isMulti()); // We only want to update one document for performance. invariant(!request.shouldReturnAnyDocs()); - invariant(PlanExecutor::NO_YIELD == request.getYieldPolicy()); + invariant(PlanYieldPolicy::YieldPolicy::NO_YIELD == request.getYieldPolicy()); auto& nss = request.getNamespaceString(); return writeConflictRetry(opCtx, "_updateWithQuery", nss.ns(), [&] { @@ -947,7 +951,7 @@ Status StorageInterfaceImpl::upsertById(OperationContext* opCtx, request.setUpsert(true); invariant(!request.isMulti()); // This follows from using an exact _id query. invariant(!request.shouldReturnAnyDocs()); - invariant(PlanExecutor::NO_YIELD == request.getYieldPolicy()); + invariant(PlanYieldPolicy::YieldPolicy::NO_YIELD == request.getYieldPolicy()); // ParsedUpdate needs to be inside the write conflict retry loop because it contains // the UpdateDriver whose state may be modified while we are applying the update. @@ -1017,7 +1021,7 @@ Status StorageInterfaceImpl::deleteByFilter(OperationContext* opCtx, request.setNsString(nss); request.setQuery(filter); request.setMulti(true); - request.setYieldPolicy(PlanExecutor::NO_YIELD); + request.setYieldPolicy(PlanYieldPolicy::YieldPolicy::NO_YIELD); // This disables the isLegalClientSystemNS() check in getExecutorDelete() which is used to // disallow client deletes from unrecognized system collections. @@ -1072,7 +1076,7 @@ boost::optional<BSONObj> StorageInterfaceImpl::findOplogEntryLessThanOrEqualToTi InternalPlanner::collectionScan(opCtx, NamespaceString::kRsOplogNamespace.ns(), oplog, - PlanExecutor::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, InternalPlanner::BACKWARD); // A record id in the oplog collection is equivalent to the document's timestamp field. diff --git a/src/mongo/db/s/migration_chunk_cloner_source_legacy.cpp b/src/mongo/db/s/migration_chunk_cloner_source_legacy.cpp index 2c311d481b9..be806add90d 100644 --- a/src/mongo/db/s/migration_chunk_cloner_source_legacy.cpp +++ b/src/mongo/db/s/migration_chunk_cloner_source_legacy.cpp @@ -829,7 +829,7 @@ MigrationChunkClonerSourceLegacy::_getIndexScanExecutor(OperationContext* opCtx, min, max, BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::YIELD_AUTO); + PlanYieldPolicy::YieldPolicy::YIELD_AUTO); } Status MigrationChunkClonerSourceLegacy::_storeCurrentLocs(OperationContext* opCtx) { diff --git a/src/mongo/db/s/range_deletion_util.cpp b/src/mongo/db/s/range_deletion_util.cpp index c97e3da2655..97c087b448c 100644 --- a/src/mongo/db/s/range_deletion_util.cpp +++ b/src/mongo/db/s/range_deletion_util.cpp @@ -188,7 +188,7 @@ StatusWith<int> deleteNextBatch(OperationContext* opCtx, min, max, BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::YIELD_MANUAL, + PlanYieldPolicy::YieldPolicy::YIELD_MANUAL, InternalPlanner::FORWARD); if (MONGO_unlikely(hangBeforeDoingDeletion.shouldFail())) { @@ -196,8 +196,6 @@ StatusWith<int> deleteNextBatch(OperationContext* opCtx, hangBeforeDoingDeletion.pauseWhileSet(opCtx); } - PlanYieldPolicy planYieldPolicy(exec.get(), PlanExecutor::YIELD_MANUAL); - int numDeleted = 0; do { BSONObj deletedObj; diff --git a/src/mongo/db/s/split_chunk.cpp b/src/mongo/db/s/split_chunk.cpp index 8d63add9683..847f43d7111 100644 --- a/src/mongo/db/s/split_chunk.cpp +++ b/src/mongo/db/s/split_chunk.cpp @@ -71,7 +71,7 @@ bool checkIfSingleDoc(OperationContext* opCtx, newmin, newmax, BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::NO_YIELD); + PlanYieldPolicy::YieldPolicy::NO_YIELD); // check if exactly one document found PlanExecutor::ExecState state; BSONObj obj; diff --git a/src/mongo/db/s/split_vector.cpp b/src/mongo/db/s/split_vector.cpp index 9eae8191854..10819bbcc8d 100644 --- a/src/mongo/db/s/split_vector.cpp +++ b/src/mongo/db/s/split_vector.cpp @@ -169,7 +169,7 @@ std::vector<BSONObj> splitVector(OperationContext* opCtx, minKey, maxKey, BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::YIELD_AUTO, + PlanYieldPolicy::YieldPolicy::YIELD_AUTO, InternalPlanner::FORWARD); BSONObj currKey; @@ -187,7 +187,7 @@ std::vector<BSONObj> splitVector(OperationContext* opCtx, maxKey, minKey, BoundInclusion::kIncludeEndKeyOnly, - PlanExecutor::YIELD_AUTO, + PlanYieldPolicy::YieldPolicy::YIELD_AUTO, InternalPlanner::BACKWARD); PlanExecutor::ExecState state = exec->getNext(&maxKeyInChunk, nullptr); @@ -304,7 +304,7 @@ std::vector<BSONObj> splitVector(OperationContext* opCtx, minKey, maxKey, BoundInclusion::kIncludeStartKeyOnly, - PlanExecutor::YIELD_AUTO, + PlanYieldPolicy::YieldPolicy::YIELD_AUTO, InternalPlanner::FORWARD); state = exec->getNext(&currKey, nullptr); diff --git a/src/mongo/db/ttl.cpp b/src/mongo/db/ttl.cpp index 987b66976b7..9eda88a8abe 100644 --- a/src/mongo/db/ttl.cpp +++ b/src/mongo/db/ttl.cpp @@ -375,7 +375,7 @@ private: startKey, endKey, BoundInclusion::kIncludeBothStartAndEndKeys, - PlanExecutor::YIELD_AUTO, + PlanYieldPolicy::YieldPolicy::YIELD_AUTO, direction); try { diff --git a/src/mongo/dbtests/cursor_manager_test.cpp b/src/mongo/dbtests/cursor_manager_test.cpp index 9f3a1ad96de..500ada8ea8e 100644 --- a/src/mongo/dbtests/cursor_manager_test.cpp +++ b/src/mongo/dbtests/cursor_manager_test.cpp @@ -79,7 +79,7 @@ public: std::move(workingSet), std::move(queuedDataStage), nullptr, - PlanExecutor::YieldPolicy::NO_YIELD, + PlanYieldPolicy::YieldPolicy::NO_YIELD, kTestNss)); } @@ -91,7 +91,6 @@ public: opCtx->getWriteConcern(), repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }; @@ -145,7 +144,6 @@ TEST_F(CursorManagerTest, ShouldBeAbleToKillPinnedCursor) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -175,7 +173,6 @@ TEST_F(CursorManagerTest, ShouldBeAbleToKillPinnedCursorMultiClient) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -216,7 +213,6 @@ TEST_F(CursorManagerTest, InactiveCursorShouldTimeout) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -236,7 +232,6 @@ TEST_F(CursorManagerTest, InactiveCursorShouldTimeout) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -260,7 +255,6 @@ TEST_F(CursorManagerTest, InactivePinnedCursorShouldNotTimeout) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -288,7 +282,6 @@ TEST_F(CursorManagerTest, MarkedAsKilledCursorsShouldBeDeletedOnCursorPin) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -325,7 +318,6 @@ TEST_F(CursorManagerTest, InactiveKilledCursorsShouldTimeout) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -361,7 +353,6 @@ TEST_F(CursorManagerTest, UsingACursorShouldUpdateTimeOfLastUse) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -379,7 +370,6 @@ TEST_F(CursorManagerTest, UsingACursorShouldUpdateTimeOfLastUse) { {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); @@ -420,7 +410,6 @@ TEST_F(CursorManagerTest, CursorShouldNotTimeOutUntilIdleForLongEnoughAfterBeing {}, repl::ReadConcernArgs(repl::ReadConcernLevel::kLocalReadConcern), BSONObj(), - ClientCursorParams::LockPolicy::kLocksInternally, PrivilegeVector(), false // needsMerge }); diff --git a/src/mongo/dbtests/deferred_writer.cpp b/src/mongo/dbtests/deferred_writer.cpp index 3d3f60b8444..07ae48ebd2c 100644 --- a/src/mongo/dbtests/deferred_writer.cpp +++ b/src/mongo/dbtests/deferred_writer.cpp @@ -117,8 +117,10 @@ public: AutoGetCollection agc(_opCtx.get(), kTestNamespace, MODE_IS); ASSERT_TRUE(agc.getCollection()); - auto plan = InternalPlanner::collectionScan( - _opCtx.get(), kTestNamespace.ns(), agc.getCollection(), PlanExecutor::NO_YIELD); + auto plan = InternalPlanner::collectionScan(_opCtx.get(), + kTestNamespace.ns(), + agc.getCollection(), + PlanYieldPolicy::YieldPolicy::NO_YIELD); std::vector<BSONObj> result; BSONObj i; diff --git a/src/mongo/dbtests/documentsourcetests.cpp b/src/mongo/dbtests/documentsourcetests.cpp index 2f6e421e159..a779fe0224d 100644 --- a/src/mongo/dbtests/documentsourcetests.cpp +++ b/src/mongo/dbtests/documentsourcetests.cpp @@ -97,8 +97,11 @@ protected: } auto cq = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(qr))); - auto exec = uassertStatusOK( - getExecutor(opCtx(), ctx.getCollection(), std::move(cq), PlanExecutor::NO_YIELD, 0)); + auto exec = uassertStatusOK(getExecutor(opCtx(), + ctx.getCollection(), + std::move(cq), + PlanYieldPolicy::YieldPolicy::NO_YIELD, + 0)); exec->saveState(); _source = DocumentSourceCursor::create( @@ -316,7 +319,7 @@ TEST_F(DocumentSourceCursorTest, TailableAwaitDataCursorShouldErrorAfterTimeout) std::move(workingSet), std::move(collectionScan), readLock.getCollection(), - PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT)); + PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT)); // Make a DocumentSourceCursor. ctx()->tailableMode = TailableModeEnum::kTailableAndAwaitData; @@ -357,7 +360,7 @@ TEST_F(DocumentSourceCursorTest, NonAwaitDataCursorShouldErrorAfterTimeout) { std::move(workingSet), std::move(collectionScan), readLock.getCollection(), - PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT)); + PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT)); // Make a DocumentSourceCursor. ctx()->tailableMode = TailableModeEnum::kNormal; @@ -407,7 +410,7 @@ TEST_F(DocumentSourceCursorTest, TailableAwaitDataCursorShouldErrorAfterBeingKil std::move(workingSet), std::move(collectionScan), readLock.getCollection(), - PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED)); + PlanYieldPolicy::YieldPolicy::ALWAYS_MARK_KILLED)); // Make a DocumentSourceCursor. ctx()->tailableMode = TailableModeEnum::kTailableAndAwaitData; @@ -447,7 +450,7 @@ TEST_F(DocumentSourceCursorTest, NormalCursorShouldErrorAfterBeingKilled) { std::move(workingSet), std::move(collectionScan), readLock.getCollection(), - PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED)); + PlanYieldPolicy::YieldPolicy::ALWAYS_MARK_KILLED)); // Make a DocumentSourceCursor. ctx()->tailableMode = TailableModeEnum::kNormal; diff --git a/src/mongo/dbtests/plan_executor_invalidation_test.cpp b/src/mongo/dbtests/plan_executor_invalidation_test.cpp index 23acefa8019..1bf6edc5a93 100644 --- a/src/mongo/dbtests/plan_executor_invalidation_test.cpp +++ b/src/mongo/dbtests/plan_executor_invalidation_test.cpp @@ -92,7 +92,7 @@ public: std::move(ws), std::move(scan), CollectionCatalog::get(&_opCtx).lookupCollectionByNamespace(&_opCtx, nss), - PlanExecutor::YIELD_MANUAL); + PlanYieldPolicy::YieldPolicy::YIELD_MANUAL); ASSERT_OK(statusWithPlanExecutor.getStatus()); return std::move(statusWithPlanExecutor.getValue()); @@ -110,7 +110,7 @@ public: startKey, endKey, BoundInclusion::kIncludeBothStartAndEndKeys, - PlanExecutor::YIELD_MANUAL); + PlanYieldPolicy::YieldPolicy::YIELD_MANUAL); } int N() { diff --git a/src/mongo/dbtests/plan_ranking.cpp b/src/mongo/dbtests/plan_ranking.cpp index 8831ee951a0..6f7991e5977 100644 --- a/src/mongo/dbtests/plan_ranking.cpp +++ b/src/mongo/dbtests/plan_ranking.cpp @@ -43,15 +43,17 @@ #include "mongo/db/db_raii.h" #include "mongo/db/dbdirectclient.h" #include "mongo/db/exec/multi_plan.h" +#include "mongo/db/exec/trial_period_utils.h" #include "mongo/db/index/index_descriptor.h" #include "mongo/db/json.h" #include "mongo/db/namespace_string.h" #include "mongo/db/query/collection_query_info.h" #include "mongo/db/query/get_executor.h" +#include "mongo/db/query/mock_yield_policies.h" #include "mongo/db/query/query_knobs_gen.h" #include "mongo/db/query/query_planner.h" #include "mongo/db/query/query_planner_test_lib.h" -#include "mongo/db/query/stage_builder.h" +#include "mongo/db/query/stage_builder_util.h" #include "mongo/dbtests/dbtests.h" namespace mongo { @@ -111,7 +113,7 @@ public: * * Does NOT take ownership of 'cq'. Caller DOES NOT own the returned QuerySolution*. */ - QuerySolution* pickBestPlan(CanonicalQuery* cq) { + const QuerySolution* pickBestPlan(CanonicalQuery* cq) { AutoGetCollectionForReadCommand ctx(&_opCtx, nss); Collection* collection = ctx.getCollection(); @@ -130,12 +132,12 @@ public: unique_ptr<WorkingSet> ws(new WorkingSet()); // Put each solution from the planner into the MPR. for (size_t i = 0; i < solutions.size(); ++i) { - auto root = StageBuilder::build(&_opCtx, collection, *cq, *solutions[i], ws.get()); + auto&& root = stage_builder::buildClassicExecutableTree( + &_opCtx, collection, *cq, *solutions[i], ws.get()); _mps->addPlan(std::move(solutions[i]), std::move(root), ws.get()); } // This is what sets a backup plan, should we test for it. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, - _opCtx.getServiceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(_opCtx.getServiceContext()->getFastClockSource()); _mps->pickBestPlan(&yieldPolicy).transitional_ignore(); ASSERT(_mps->bestPlanChosen()); @@ -143,7 +145,7 @@ public: ASSERT_LESS_THAN(bestPlanIdx, solutions.size()); // And return a pointer to the best solution. - return _mps->bestSolution(); + return static_cast<const MultiPlanStage*>(_mps.get())->bestSolution(); } /** @@ -212,7 +214,7 @@ public: // We get the number of works done during the trial period in order to make sure that there // are more documents in the collection than works done in the trial period. This ensures // neither of the plans reach EOF or produce results. - size_t numWorks = MultiPlanStage::getTrialPeriodWorks(opCtx(), nullptr); + size_t numWorks = trial_period::getTrialPeriodMaxWorks(opCtx(), nullptr); size_t smallNumber = 10; // The following condition must be met in order for the following test to work. Specifically // this condition guarantees that the score of the plan using the index on d will score @@ -237,7 +239,7 @@ public: unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); ASSERT(cq); - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT( QueryPlannerTestLib::solutionMatches("{fetch: {filter: {a:1}, node: " "{ixscan: {filter: null, pattern: {d:1}}}}}", @@ -292,7 +294,7 @@ public: } // {a:100} is super selective so choose that. - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT(QueryPlannerTestLib::solutionMatches( "{fetch: {filter: {b:1}, node: {ixscan: {pattern: {a: 1}}}}}", soln->root.get())); @@ -348,7 +350,7 @@ public: // This will be reverted by PlanRankingTestBase's destructor when the test completes. internalQueryForceIntersectionPlans.store(true); - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT( QueryPlannerTestLib::solutionMatches("{fetch: {node: {andHash: {nodes: [" "{ixscan: {filter: null, pattern: {a:1}}}," @@ -385,7 +387,7 @@ public: unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); ASSERT(nullptr != cq.get()); - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); // Prefer the fully covered plan. ASSERT(QueryPlannerTestLib::solutionMatches( @@ -419,7 +421,7 @@ public: unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); ASSERT(nullptr != cq.get()); - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); // Anti-prefer the intersection plan. bool bestIsScanOverA = QueryPlannerTestLib::solutionMatches( @@ -458,7 +460,7 @@ public: unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); ASSERT(nullptr != cq.get()); - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); // Prefer the fully covered plan. ASSERT(QueryPlannerTestLib::solutionMatches( "{proj: {spec: {_id:0, a:1, b:1}, node: {ixscan: {pattern: {a: 1, b:1}}}}}", @@ -492,7 +494,7 @@ public: ASSERT(nullptr != cq.get()); // {a: 100} is super selective so choose that. - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT(QueryPlannerTestLib::solutionMatches( "{fetch: {filter: {b:1}, node: {ixscan: {pattern: {a: 1}}}}}", soln->root.get())); } @@ -528,7 +530,7 @@ public: ASSERT(nullptr != cq.get()); // {a: 100} is super selective so choose that. - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT(QueryPlannerTestLib::solutionMatches( "{fetch: {filter: {b:1}, node: {ixscan: {pattern: {a: 1}}}}}", soln->root.get())); } @@ -556,7 +558,7 @@ public: ASSERT_OK(statusWithCQ.getStatus()); unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); // The best must not be a collscan. ASSERT(QueryPlannerTestLib::solutionMatches( @@ -586,7 +588,7 @@ public: unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); ASSERT(nullptr != cq.get()); - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); // The best must be a collscan. ASSERT(QueryPlannerTestLib::solutionMatches("{cscan: {dir: 1, filter: {foo: 2001}}}", @@ -623,7 +625,7 @@ public: // No results will be returned during the trial period, // so we expect to choose {d: 1, e: 1}, as it allows us // to avoid the sort stage. - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT( QueryPlannerTestLib::solutionMatches("{fetch: {filter: {a:1}, node: " "{ixscan: {filter: null, pattern: {d:1,e:1}}}}}", @@ -658,7 +660,7 @@ public: ASSERT(nullptr != cq.get()); // Use index on 'b'. - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT(QueryPlannerTestLib::solutionMatches("{fetch: {node: {ixscan: {pattern: {b: 1}}}}}", soln->root.get())); } @@ -690,7 +692,7 @@ public: ASSERT(nullptr != cq.get()); // Expect to use index {a: 1, b: 1}. - QuerySolution* soln = pickBestPlan(cq.get()); + auto soln = pickBestPlan(cq.get()); ASSERT(QueryPlannerTestLib::solutionMatches("{fetch: {node: {ixscan: {pattern: {a: 1}}}}}", soln->root.get())); } diff --git a/src/mongo/dbtests/query_plan_executor.cpp b/src/mongo/dbtests/query_plan_executor.cpp index 529a572097f..10961cb1018 100644 --- a/src/mongo/dbtests/query_plan_executor.cpp +++ b/src/mongo/dbtests/query_plan_executor.cpp @@ -99,7 +99,7 @@ public: unique_ptr<PlanExecutor, PlanExecutor::Deleter> makeCollScanExec( Collection* coll, BSONObj& filterObj, - PlanExecutor::YieldPolicy yieldPolicy = PlanExecutor::YieldPolicy::YIELD_MANUAL, + PlanYieldPolicy::YieldPolicy yieldPolicy = PlanYieldPolicy::YieldPolicy::YIELD_MANUAL, TailableModeEnum tailableMode = TailableModeEnum::kNormal) { CollectionScanParams csparams; csparams.direction = CollectionScanParams::FORWARD; @@ -164,8 +164,12 @@ public: verify(nullptr != cq.get()); // Hand the plan off to the executor. - auto statusWithPlanExecutor = PlanExecutor::make( - std::move(cq), std::move(ws), std::move(root), coll, PlanExecutor::YIELD_MANUAL); + auto statusWithPlanExecutor = + PlanExecutor::make(std::move(cq), + std::move(ws), + std::move(root), + coll, + PlanYieldPolicy::YieldPolicy::YIELD_MANUAL); ASSERT_OK(statusWithPlanExecutor.getStatus()); return std::move(statusWithPlanExecutor.getValue()); } @@ -223,8 +227,11 @@ TEST_F(PlanExecutorTest, DropIndexScanAgg) { auto ws = std::make_unique<WorkingSet>(); auto proxy = std::make_unique<PipelineProxyStage>(_expCtx.get(), std::move(pipeline), ws.get()); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(proxy), collection, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(proxy), + collection, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto outerExec = std::move(statusWithPlanExecutor.getValue()); @@ -245,7 +252,7 @@ TEST_F(PlanExecutorTest, ShouldReportErrorIfExceedsTimeLimitDuringYield) { BSONObj filterObj = fromjson("{_id: {$gt: 0}}"); Collection* coll = ctx.getCollection(); - auto exec = makeCollScanExec(coll, filterObj, PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT); + auto exec = makeCollScanExec(coll, filterObj, PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT); BSONObj resultObj; ASSERT_THROWS_CODE_AND_WHAT(exec->getNext(&resultObj, nullptr), @@ -264,7 +271,7 @@ TEST_F(PlanExecutorTest, ShouldReportErrorIfKilledDuringYieldButIsTailableAndAwa Collection* coll = ctx.getCollection(); auto exec = makeCollScanExec(coll, filterObj, - PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT, + PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT, TailableModeEnum::kTailableAndAwaitData); BSONObj resultObj; @@ -282,8 +289,10 @@ TEST_F(PlanExecutorTest, ShouldNotSwallowExceedsTimeLimitDuringYieldButIsTailabl BSONObj filterObj = fromjson("{_id: {$gt: 0}}"); Collection* coll = ctx.getCollection(); - auto exec = makeCollScanExec( - coll, filterObj, PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT, TailableModeEnum::kTailable); + auto exec = makeCollScanExec(coll, + filterObj, + PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT, + TailableModeEnum::kTailable); BSONObj resultObj; ASSERT_THROWS_CODE_AND_WHAT(exec->getNext(&resultObj, nullptr), @@ -300,7 +309,7 @@ TEST_F(PlanExecutorTest, ShouldReportErrorIfKilledDuringYield) { BSONObj filterObj = fromjson("{_id: {$gt: 0}}"); Collection* coll = ctx.getCollection(); - auto exec = makeCollScanExec(coll, filterObj, PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED); + auto exec = makeCollScanExec(coll, filterObj, PlanYieldPolicy::YieldPolicy::ALWAYS_MARK_KILLED); BSONObj resultObj; ASSERT_THROWS_CODE_AND_WHAT(exec->getNext(&resultObj, nullptr), diff --git a/src/mongo/dbtests/query_stage_cached_plan.cpp b/src/mongo/dbtests/query_stage_cached_plan.cpp index 87ea3b7e3bd..ba32d9235d6 100644 --- a/src/mongo/dbtests/query_stage_cached_plan.cpp +++ b/src/mongo/dbtests/query_stage_cached_plan.cpp @@ -46,8 +46,8 @@ #include "mongo/db/query/canonical_query.h" #include "mongo/db/query/collection_query_info.h" #include "mongo/db/query/get_executor.h" +#include "mongo/db/query/mock_yield_policies.h" #include "mongo/db/query/plan_cache.h" -#include "mongo/db/query/plan_yield_policy.h" #include "mongo/db/query/query_knobs_gen.h" #include "mongo/db/query/query_planner_params.h" #include "mongo/dbtests/dbtests.h" @@ -164,8 +164,7 @@ public: std::move(mockChild)); // This should succeed after triggering a replan. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, - _opCtx.getServiceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(_opCtx.getServiceContext()->getFastClockSource()); ASSERT_OK(cachedPlanStage.pickBestPlan(&yieldPolicy)); } @@ -220,8 +219,7 @@ TEST_F(QueryStageCachedPlan, QueryStageCachedPlanFailureMemoryLimitExceeded) { std::move(mockChild)); // This should succeed after triggering a replan. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, - _opCtx.getServiceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(_opCtx.getServiceContext()->getFastClockSource()); ASSERT_OK(cachedPlanStage.pickBestPlan(&yieldPolicy)); ASSERT_EQ(getNumResultsForStage(_ws, &cachedPlanStage, cq.get()), 2U); @@ -275,8 +273,7 @@ TEST_F(QueryStageCachedPlan, QueryStageCachedPlanHitMaxWorks) { std::move(mockChild)); // This should succeed after triggering a replan. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, - _opCtx.getServiceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(_opCtx.getServiceContext()->getFastClockSource()); ASSERT_OK(cachedPlanStage.pickBestPlan(&yieldPolicy)); ASSERT_EQ(getNumResultsForStage(_ws, &cachedPlanStage, cq.get()), 2U); @@ -520,8 +517,7 @@ TEST_F(QueryStageCachedPlan, DoesNotThrowOnYieldRecoveryWhenIndexIsDroppedAferPl decisionWorks, std::make_unique<MockStage>(_expCtx.get(), &_ws)); - PlanYieldPolicy yieldPolicy(PlanExecutor::YIELD_MANUAL, - _opCtx.getServiceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(_opCtx.getServiceContext()->getFastClockSource()); ASSERT_OK(cachedPlanStage.pickBestPlan(&yieldPolicy)); // Drop an index while the CachedPlanStage is in a saved state. We should be able to restore diff --git a/src/mongo/dbtests/query_stage_collscan.cpp b/src/mongo/dbtests/query_stage_collscan.cpp index c9b026119bf..27fd43dd268 100644 --- a/src/mongo/dbtests/query_stage_collscan.cpp +++ b/src/mongo/dbtests/query_stage_collscan.cpp @@ -104,8 +104,11 @@ public: unique_ptr<PlanStage> ps = std::make_unique<CollectionScan>( _expCtx.get(), collection, params, ws.get(), filterExpr.get()); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(ps), collection, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(ps), + collection, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -195,7 +198,7 @@ TEST_F(QueryStageCollectionScanTest, QueryStageCollscanObjectsInOrderForward) { std::make_unique<CollectionScan>(_expCtx.get(), collection, params, ws.get(), nullptr); auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(ps), collection, PlanExecutor::NO_YIELD); + _expCtx, std::move(ws), std::move(ps), collection, PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -224,7 +227,7 @@ TEST_F(QueryStageCollectionScanTest, QueryStageCollscanObjectsInOrderBackward) { std::make_unique<CollectionScan>(_expCtx.get(), collection, params, ws.get(), nullptr); auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(ps), collection, PlanExecutor::NO_YIELD); + _expCtx, std::move(ws), std::move(ps), collection, PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -377,7 +380,7 @@ TEST_F(QueryStageCollectionScanTest, QueryTestCollscanResumeAfterRecordIdSeekSuc // Run the rest of the scan and verify the results. auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(ps), collection, PlanExecutor::NO_YIELD); + _expCtx, std::move(ws), std::move(ps), collection, PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); diff --git a/src/mongo/dbtests/query_stage_merge_sort.cpp b/src/mongo/dbtests/query_stage_merge_sort.cpp index 3ad17b1cf52..ec91fcc4607 100644 --- a/src/mongo/dbtests/query_stage_merge_sort.cpp +++ b/src/mongo/dbtests/query_stage_merge_sort.cpp @@ -186,8 +186,11 @@ public: unique_ptr<FetchStage> fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); // Must fetch if we want to easily pull out an obj. - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -251,8 +254,11 @@ public: unique_ptr<FetchStage> fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -316,8 +322,11 @@ public: unique_ptr<FetchStage> fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -387,8 +396,11 @@ public: unique_ptr<FetchStage> fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -454,8 +466,11 @@ public: unique_ptr<FetchStage> fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -508,8 +523,11 @@ public: unique_ptr<FetchStage> fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -797,8 +815,11 @@ public: auto fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); // Must fetch if we want to easily pull out an obj. - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -866,8 +887,11 @@ public: unique_ptr<FetchStage> fetchStage = make_unique<FetchStage>(_expCtx.get(), ws.get(), std::move(ms), nullptr, coll); // Must fetch if we want to easily pull out an obj. - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); diff --git a/src/mongo/dbtests/query_stage_multiplan.cpp b/src/mongo/dbtests/query_stage_multiplan.cpp index 1e12492f4e5..ca75cc3052b 100644 --- a/src/mongo/dbtests/query_stage_multiplan.cpp +++ b/src/mongo/dbtests/query_stage_multiplan.cpp @@ -54,7 +54,7 @@ #include "mongo/db/query/query_knobs_gen.h" #include "mongo/db/query/query_planner.h" #include "mongo/db/query/query_planner_test_lib.h" -#include "mongo/db/query/stage_builder.h" +#include "mongo/db/query/stage_builder_util.h" #include "mongo/dbtests/dbtests.h" #include "mongo/util/clock_source_mock.h" @@ -202,8 +202,7 @@ std::unique_ptr<MultiPlanStage> runMultiPlanner(ExpressionContext* expCtx, mps->addPlan(createQuerySolution(), std::move(collScanRoot), sharedWs.get()); // Plan 0 aka the first plan aka the index scan should be the best. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, - expCtx->opCtx->getServiceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(expCtx->opCtx->getServiceContext()->getFastClockSource()); ASSERT_OK(mps->pickBestPlan(&yieldPolicy)); ASSERT(mps->bestPlanChosen()); ASSERT_EQUALS(0, mps->bestPlanIdx()); @@ -250,14 +249,17 @@ TEST_F(QueryStageMultiPlanTest, MPSCollectionScanVsHighlySelectiveIXScan) { mps->addPlan(createQuerySolution(), std::move(collScanRoot), sharedWs.get()); // Plan 0 aka the first plan aka the index scan should be the best. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock); + NoopYieldPolicy yieldPolicy(_clock); ASSERT_OK(mps->pickBestPlan(&yieldPolicy)); ASSERT(mps->bestPlanChosen()); ASSERT_EQUALS(0, mps->bestPlanIdx()); // Takes ownership of arguments other than 'collection'. - auto statusWithPlanExecutor = PlanExecutor::make( - std::move(cq), std::move(sharedWs), std::move(mps), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(std::move(cq), + std::move(sharedWs), + std::move(mps), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -400,18 +402,19 @@ TEST_F(QueryStageMultiPlanTest, MPSBackupPlan) { unique_ptr<WorkingSet> ws(new WorkingSet()); // Put each solution from the planner into the MPR. for (size_t i = 0; i < solutions.size(); ++i) { - auto root = StageBuilder::build(_opCtx.get(), collection, *cq, *solutions[i], ws.get()); + auto&& root = stage_builder::buildClassicExecutableTree( + _opCtx.get(), collection, *cq, *solutions[i], ws.get()); mps->addPlan(std::move(solutions[i]), std::move(root), ws.get()); } // This sets a backup plan. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock); + NoopYieldPolicy yieldPolicy(_clock); ASSERT_OK(mps->pickBestPlan(&yieldPolicy)); ASSERT(mps->bestPlanChosen()); ASSERT(mps->hasBackupPlan()); // We should have picked the index intersection plan due to forcing ixisect. - QuerySolution* soln = mps->bestSolution(); + auto soln = static_cast<const MultiPlanStage*>(mps.get())->bestSolution(); ASSERT( QueryPlannerTestLib::solutionMatches("{sort: {pattern: {b: 1}, limit: 0, node:" "{fetch: {node: {andSorted: {nodes: [" @@ -435,7 +438,7 @@ TEST_F(QueryStageMultiPlanTest, MPSBackupPlan) { // The blocking plan became unblocked, so we should no longer have a backup plan, // and the winning plan should still be the index intersection one. ASSERT(!mps->hasBackupPlan()); - soln = mps->bestSolution(); + soln = static_cast<const MultiPlanStage*>(mps.get())->bestSolution(); ASSERT( QueryPlannerTestLib::solutionMatches("{sort: {pattern: {b: 1}, limit: 0, node:" "{fetch: {node: {andSorted: {nodes: [" @@ -491,8 +494,11 @@ TEST_F(QueryStageMultiPlanTest, MPSExplainAllPlans) { mps->addPlan(std::make_unique<QuerySolution>(), std::move(secondPlan), ws.get()); // Making a PlanExecutor chooses the best plan. - auto exec = uassertStatusOK(PlanExecutor::make( - _expCtx, std::move(ws), std::move(mps), ctx.getCollection(), PlanExecutor::NO_YIELD)); + auto exec = uassertStatusOK(PlanExecutor::make(_expCtx, + std::move(ws), + std::move(mps), + ctx.getCollection(), + PlanYieldPolicy::YieldPolicy::NO_YIELD)); auto root = static_cast<MultiPlanStage*>(exec->getRootStage()); ASSERT_TRUE(root->bestPlanChosen()); @@ -541,8 +547,8 @@ TEST_F(QueryStageMultiPlanTest, MPSSummaryStats) { auto qr = std::make_unique<QueryRequest>(nss); qr->setFilter(BSON("foo" << BSON("$gte" << 0))); auto cq = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(qr))); - auto exec = - uassertStatusOK(getExecutor(opCtx(), coll, std::move(cq), PlanExecutor::NO_YIELD, 0)); + auto exec = uassertStatusOK( + getExecutor(opCtx(), coll, std::move(cq), PlanYieldPolicy::YieldPolicy::NO_YIELD, 0)); ASSERT_EQ(exec->getRootStage()->stageType(), STAGE_MULTI_PLAN); exec->executePlan(); @@ -586,10 +592,8 @@ TEST_F(QueryStageMultiPlanTest, ShouldReportErrorIfExceedsTimeLimitDuringPlannin queryRequest->setFilter(filterObj); auto canonicalQuery = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest))); - MultiPlanStage multiPlanStage(_expCtx.get(), - ctx.getCollection(), - canonicalQuery.get(), - MultiPlanStage::CachingMode::NeverCache); + MultiPlanStage multiPlanStage( + _expCtx.get(), ctx.getCollection(), canonicalQuery.get(), PlanCachingMode::NeverCache); multiPlanStage.addPlan(createQuerySolution(), std::move(ixScanRoot), sharedWs.get()); multiPlanStage.addPlan(createQuerySolution(), std::move(collScanRoot), sharedWs.get()); @@ -628,10 +632,8 @@ TEST_F(QueryStageMultiPlanTest, ShouldReportErrorIfKilledDuringPlanning) { queryRequest->setFilter(BSON("foo" << BSON("$gte" << 0))); auto canonicalQuery = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest))); - MultiPlanStage multiPlanStage(_expCtx.get(), - ctx.getCollection(), - canonicalQuery.get(), - MultiPlanStage::CachingMode::NeverCache); + MultiPlanStage multiPlanStage( + _expCtx.get(), ctx.getCollection(), canonicalQuery.get(), PlanCachingMode::NeverCache); multiPlanStage.addPlan(createQuerySolution(), std::move(ixScanRoot), sharedWs.get()); multiPlanStage.addPlan(createQuerySolution(), std::move(collScanRoot), sharedWs.get()); @@ -674,17 +676,15 @@ TEST_F(QueryStageMultiPlanTest, AddsContextDuringException) { << "query")); auto canonicalQuery = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest))); - MultiPlanStage multiPlanStage(_expCtx.get(), - ctx.getCollection(), - canonicalQuery.get(), - MultiPlanStage::CachingMode::NeverCache); + MultiPlanStage multiPlanStage( + _expCtx.get(), ctx.getCollection(), canonicalQuery.get(), PlanCachingMode::NeverCache); unique_ptr<WorkingSet> sharedWs(new WorkingSet()); multiPlanStage.addPlan( createQuerySolution(), std::make_unique<ThrowyPlanStage>(_expCtx.get()), sharedWs.get()); multiPlanStage.addPlan( createQuerySolution(), std::make_unique<ThrowyPlanStage>(_expCtx.get()), sharedWs.get()); - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock); + NoopYieldPolicy yieldPolicy(_clock); auto status = multiPlanStage.pickBestPlan(&yieldPolicy); ASSERT_EQ(ErrorCodes::InternalError, status); ASSERT_STRING_CONTAINS(status.reason(), "error while multiplanner was selecting best plan"); diff --git a/src/mongo/dbtests/query_stage_sort.cpp b/src/mongo/dbtests/query_stage_sort.cpp index dc49b5c49f5..612771d6425 100644 --- a/src/mongo/dbtests/query_stage_sort.cpp +++ b/src/mongo/dbtests/query_stage_sort.cpp @@ -131,7 +131,7 @@ public: // The PlanExecutor will be automatically registered on construction due to the auto // yield policy, so it can receive invalidations when we remove documents later. auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(ss), coll, PlanExecutor::YIELD_AUTO); + _expCtx, std::move(ws), std::move(ss), coll, PlanYieldPolicy::YieldPolicy::YIELD_AUTO); invariant(statusWithPlanExecutor.isOK()); return std::move(statusWithPlanExecutor.getValue()); } @@ -173,8 +173,11 @@ public: _expCtx.get(), ws.get(), std::move(sortStage), nullptr, coll); // Must fetch so we can look at the doc as a BSONObj. - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); @@ -587,8 +590,11 @@ public: _expCtx.get(), ws.get(), std::move(sortStage), nullptr, coll); // We don't get results back since we're sorting some parallel arrays. - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(fetchStage), coll, PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(fetchStage), + coll, + PlanYieldPolicy::YieldPolicy::NO_YIELD); auto exec = std::move(statusWithPlanExecutor.getValue()); ASSERT_THROWS_CODE(exec->getNext(static_cast<BSONObj*>(nullptr), nullptr), diff --git a/src/mongo/dbtests/query_stage_subplan.cpp b/src/mongo/dbtests/query_stage_subplan.cpp index f1b2ec000df..8d56d4e094d 100644 --- a/src/mongo/dbtests/query_stage_subplan.cpp +++ b/src/mongo/dbtests/query_stage_subplan.cpp @@ -150,7 +150,7 @@ TEST_F(QueryStageSubplanTest, QueryStageSubplanGeo2dOr) { new SubplanStage(_expCtx.get(), collection, &ws, plannerParams, cq.get())); // Plan selection should succeed due to falling back on regular planning. - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock); + NoopYieldPolicy yieldPolicy(_clock); ASSERT_OK(subplan->pickBestPlan(&yieldPolicy)); } @@ -188,8 +188,7 @@ void assertSubplanFromCache(QueryStageSubplanTest* test, const dbtests::WriteCon std::unique_ptr<SubplanStage> subplan( new SubplanStage(test->expCtx(), collection, &ws, plannerParams, cq.get())); - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, - test->serviceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(test->serviceContext()->getFastClockSource()); ASSERT_OK(subplan->pickBestPlan(&yieldPolicy)); // Nothing is in the cache yet, so neither branch should have been planned from @@ -269,7 +268,7 @@ TEST_F(QueryStageSubplanTest, QueryStageSubplanDontCacheZeroResults) { std::unique_ptr<SubplanStage> subplan( new SubplanStage(_expCtx.get(), collection, &ws, plannerParams, cq.get())); - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock); + NoopYieldPolicy yieldPolicy(_clock); ASSERT_OK(subplan->pickBestPlan(&yieldPolicy)); // Nothing is in the cache yet, so neither branch should have been planned from @@ -325,7 +324,7 @@ TEST_F(QueryStageSubplanTest, QueryStageSubplanDontCacheTies) { std::unique_ptr<SubplanStage> subplan( new SubplanStage(_expCtx.get(), collection, &ws, plannerParams, cq.get())); - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock); + NoopYieldPolicy yieldPolicy(_clock); ASSERT_OK(subplan->pickBestPlan(&yieldPolicy)); // Nothing is in the cache yet, so neither branch should have been planned from @@ -497,7 +496,7 @@ TEST_F(QueryStageSubplanTest, QueryStageSubplanPlanRootedOrNE) { std::unique_ptr<SubplanStage> subplan( new SubplanStage(_expCtx.get(), collection, &ws, plannerParams, cq.get())); - PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock); + NoopYieldPolicy yieldPolicy(_clock); ASSERT_OK(subplan->pickBestPlan(&yieldPolicy)); size_t numResults = 0; @@ -647,7 +646,7 @@ TEST_F(QueryStageSubplanTest, ShouldNotThrowOnRestoreIfIndexDroppedAfterPlanSele WorkingSet workingSet; SubplanStage subplanStage(_expCtx.get(), collection, &workingSet, params, canonicalQuery.get()); - PlanYieldPolicy yieldPolicy(PlanExecutor::YIELD_MANUAL, serviceContext()->getFastClockSource()); + NoopYieldPolicy yieldPolicy(serviceContext()->getFastClockSource()); ASSERT_OK(subplanStage.pickBestPlan(&yieldPolicy)); // Mimic a yield by saving the state of the subplan stage and dropping our lock. Then drop an diff --git a/src/mongo/dbtests/query_stage_tests.cpp b/src/mongo/dbtests/query_stage_tests.cpp index 1c705c3ff12..e51f95f830b 100644 --- a/src/mongo/dbtests/query_stage_tests.cpp +++ b/src/mongo/dbtests/query_stage_tests.cpp @@ -92,8 +92,11 @@ public: unique_ptr<IndexScan> ix = std::make_unique<IndexScan>( _expCtx.get(), ctx.getCollection(), params, ws.get(), filterExpr.get()); - auto statusWithPlanExecutor = PlanExecutor::make( - _expCtx, std::move(ws), std::move(ix), ctx.getCollection(), PlanExecutor::NO_YIELD); + auto statusWithPlanExecutor = PlanExecutor::make(_expCtx, + std::move(ws), + std::move(ix), + ctx.getCollection(), + PlanYieldPolicy::YieldPolicy::NO_YIELD); ASSERT_OK(statusWithPlanExecutor.getStatus()); auto exec = std::move(statusWithPlanExecutor.getValue()); diff --git a/src/mongo/dbtests/query_stage_trial.cpp b/src/mongo/dbtests/query_stage_trial.cpp index 72bb16ba34f..00f7bdbcfa9 100644 --- a/src/mongo/dbtests/query_stage_trial.cpp +++ b/src/mongo/dbtests/query_stage_trial.cpp @@ -36,7 +36,7 @@ #include "mongo/db/exec/trial_stage.h" #include "mongo/db/exec/working_set.h" #include "mongo/db/operation_context.h" -#include "mongo/db/query/plan_yield_policy.h" +#include "mongo/db/query/mock_yield_policies.h" #include "mongo/db/service_context_d_test_fixture.h" #include "mongo/unittest/unittest.h" @@ -87,8 +87,8 @@ protected: } std::unique_ptr<PlanYieldPolicy> yieldPolicy() { - return std::make_unique<PlanYieldPolicy>( - PlanExecutor::NO_YIELD, opCtx()->getServiceContext()->getFastClockSource()); + return std::make_unique<NoopYieldPolicy>( + opCtx()->getServiceContext()->getFastClockSource()); } OperationContext* opCtx() { diff --git a/src/mongo/platform/decimal128.h b/src/mongo/platform/decimal128.h index 54e169330a9..95215dea62b 100644 --- a/src/mongo/platform/decimal128.h +++ b/src/mongo/platform/decimal128.h @@ -605,6 +605,30 @@ inline Decimal128 operator"" _dec128(const char* s, std::size_t len) { return Decimal128(std::string(s, len)); } +inline bool operator<(const Decimal128& lhs, const Decimal128& rhs) { + return lhs.isLess(rhs); +} + +inline bool operator<=(const Decimal128& lhs, const Decimal128& rhs) { + return lhs.isLessEqual(rhs); +} + +inline bool operator>(const Decimal128& lhs, const Decimal128& rhs) { + return lhs.isGreater(rhs); +} + +inline bool operator>=(const Decimal128& lhs, const Decimal128& rhs) { + return lhs.isGreaterEqual(rhs); +} + +inline bool operator==(const Decimal128& lhs, const Decimal128& rhs) { + return lhs.isEqual(rhs); +} + +inline bool operator!=(const Decimal128& lhs, const Decimal128& rhs) { + return lhs.isNotEqual(rhs); +} + } // namespace literals } // namespace mongo |