diff options
Diffstat (limited to 'src/mongo/db')
21 files changed, 360 insertions, 145 deletions
diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript index fef38386595..2112af67b86 100644 --- a/src/mongo/db/SConscript +++ b/src/mongo/db/SConscript @@ -991,6 +991,7 @@ env.Library( 'exec/requires_all_indices_stage.cpp', 'exec/requires_collection_stage.cpp', 'exec/requires_index_stage.cpp', + 'exec/return_key.cpp', 'exec/shard_filter.cpp', 'exec/shard_filterer_impl.cpp', 'exec/skip.cpp', diff --git a/src/mongo/db/exec/idhack.cpp b/src/mongo/db/exec/idhack.cpp index 4ff011d024e..02e347dda4d 100644 --- a/src/mongo/db/exec/idhack.cpp +++ b/src/mongo/db/exec/idhack.cpp @@ -57,9 +57,7 @@ IDHackStage::IDHackStage(OperationContext* opCtx, _workingSet(ws), _key(query->getQueryObj()["_id"].wrap()) { _specificStats.indexName = descriptor->indexName(); - if (nullptr != query->getProj()) { - _addKeyMetadata = query->getProj()->wantIndexKey(); - } + _addKeyMetadata = query->getQueryRequest().returnKey(); } IDHackStage::IDHackStage(OperationContext* opCtx, diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h index 9ec4d0b69a5..cff5085e47a 100644 --- a/src/mongo/db/exec/plan_stats.h +++ b/src/mongo/db/exec/plan_stats.h @@ -447,6 +447,16 @@ struct IDHackStats : public SpecificStats { size_t docsExamined; }; +struct ReturnKeyStats : public SpecificStats { + SpecificStats* clone() const final { + return new ReturnKeyStats(*this); + } + + uint64_t estimateObjectSizeInBytes() const { + return sizeof(*this); + } +}; + struct IndexScanStats : public SpecificStats { IndexScanStats() : indexVersion(0), diff --git a/src/mongo/db/exec/projection.cpp b/src/mongo/db/exec/projection.cpp index cc7ba51c095..8faa479e9d9 100644 --- a/src/mongo/db/exec/projection.cpp +++ b/src/mongo/db/exec/projection.cpp @@ -205,19 +205,7 @@ Status ProjectionStageDefault::transform(WorkingSetMember* member) const { return Status(ErrorCodes::InternalError, "sortKey meta-projection requested but no data available"); - if (_exec.returnKey()) { - auto keys = _exec.computeReturnKeyProjection( - member->metadata().hasIndexKey() ? indexKey(*member) : BSONObj(), - _exec.needsSortKey() ? sortKey(*member) : BSONObj()); - if (!keys.isOK()) - return keys.getStatus(); - - transitionMemberToOwnedObj(keys.getValue(), member); - return Status::OK(); - } - auto projected = provideMetaFieldsAndPerformExec(_exec, *member); - if (!projected.isOK()) return projected.getStatus(); diff --git a/src/mongo/db/exec/projection_exec.cpp b/src/mongo/db/exec/projection_exec.cpp index 736dd900545..bf212a0b44b 100644 --- a/src/mongo/db/exec/projection_exec.cpp +++ b/src/mongo/db/exec/projection_exec.cpp @@ -106,8 +106,7 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx, _meta[e.fieldName()] = META_TEXT_SCORE; _needsTextScore = true; } else if (e2.valuestr() == QueryRequest::metaSortKey) { - _sortKeyMetaFields.push_back(e.fieldName()); - _meta[_sortKeyMetaFields.back()] = META_SORT_KEY; + _meta[e.fieldName()] = META_SORT_KEY; _needsSortKey = true; } else if (e2.valuestr() == QueryRequest::metaRecordId) { _meta[e.fieldName()] = META_RECORDID; @@ -117,8 +116,6 @@ ProjectionExec::ProjectionExec(OperationContext* opCtx, } else if (e2.valuestr() == QueryRequest::metaGeoNearDistance) { _meta[e.fieldName()] = META_GEONEAR_DIST; _needsGeoNearDistance = true; - } else if (e2.valuestr() == QueryRequest::metaIndexKey) { - _hasReturnKey = true; } else { // This shouldn't happen, should be caught by parsing. MONGO_UNREACHABLE; @@ -200,22 +197,6 @@ void ProjectionExec::add(const string& field, int skip, int limit) { // Execution // -StatusWith<BSONObj> ProjectionExec::computeReturnKeyProjection(const BSONObj& indexKey, - const BSONObj& sortKey) const { - BSONObjBuilder bob; - - if (!indexKey.isEmpty()) { - bob.appendElements(indexKey); - } - - // Must be possible to do both returnKey meta-projection and sortKey meta-projection so that - // mongos can support returnKey. - for (auto fieldName : _sortKeyMetaFields) - bob.append(fieldName, sortKey); - - return bob.obj(); -} - StatusWith<BSONObj> ProjectionExec::project(const BSONObj& in, const boost::optional<const double> geoDistance, Value geoNearPoint, diff --git a/src/mongo/db/exec/projection_exec.h b/src/mongo/db/exec/projection_exec.h index a6ce7b1317c..c1b878aa48c 100644 --- a/src/mongo/db/exec/projection_exec.h +++ b/src/mongo/db/exec/projection_exec.h @@ -80,14 +80,6 @@ public: ~ProjectionExec(); /** - * Indicates whether this is a returnKey projection which should be performed via - * 'computeReturnKeyProjection()'. - */ - bool returnKey() const { - return _hasReturnKey; - } - - /** * Indicates whether 'sortKey' must be provided for 'computeReturnKeyProjection()' or * 'project()'. */ @@ -124,12 +116,6 @@ public: } /** - * Performs a returnKey projection and provides index keys rather than projection results. - */ - StatusWith<BSONObj> computeReturnKeyProjection(const BSONObj& indexKey, - const BSONObj& sortKey) const; - - /** * Performs a projection given a BSONObj source. Meta fields must be provided if necessary. * Their necessity can be queried via the 'needs*' functions. */ @@ -259,10 +245,6 @@ private: // Projections that aren't sourced from the document or index keys. MetaMap _meta; - // Do we have a returnKey projection? If so we *only* output the index key metadata, and - // possibly the sort key for mongos to use. If it's not found we output nothing. - bool _hasReturnKey = false; - // After parsing in the constructor, these fields will indicate the neccesity of metadata // for $meta projection. bool _needsSortKey = false; @@ -270,10 +252,6 @@ private: bool _needsGeoNearPoint = false; bool _needsTextScore = false; - // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey - // meta-projection. - std::vector<StringData> _sortKeyMetaFields; - // The collator this projection should use to compare strings. Needed for projection operators // that perform matching (e.g. elemMatch projection). If null, the collation is a simple binary // compare. diff --git a/src/mongo/db/exec/return_key.cpp b/src/mongo/db/exec/return_key.cpp new file mode 100644 index 00000000000..8f0e185b554 --- /dev/null +++ b/src/mongo/db/exec/return_key.cpp @@ -0,0 +1,103 @@ +/** + * 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_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kQuery + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/return_key.h" + +#include "mongo/db/exec/working_set_common.h" +#include "mongo/util/log.h" + +namespace mongo { +using namespace fmt::literals; + +PlanStage::StageState ReturnKeyStage::doWork(WorkingSetID* out) { + WorkingSetID id = WorkingSet::INVALID_ID; + StageState status = child()->work(&id); + + if (PlanStage::ADVANCED == status) { + WorkingSetMember* member = _ws.get(id); + Status indexKeyStatus = _extractIndexKey(member); + + if (!indexKeyStatus.isOK()) { + warning() << "Couldn't execute {}, status = {}"_format(kStageName, + redact(indexKeyStatus)); + *out = WorkingSetCommon::allocateStatusMember(&_ws, indexKeyStatus); + return PlanStage::FAILURE; + } + + *out = id; + } else if (PlanStage::FAILURE == status) { + // The stage which produces a failure is responsible for allocating a working set member + // with error details. + invariant(WorkingSet::INVALID_ID != id); + *out = id; + } else if (PlanStage::NEED_YIELD == status) { + *out = id; + } + + return status; +} + +std::unique_ptr<PlanStageStats> ReturnKeyStage::getStats() { + _commonStats.isEOF = isEOF(); + + auto ret = std::make_unique<PlanStageStats>(_commonStats, stageType()); + ret->specific = std::make_unique<ReturnKeyStats>(_specificStats); + ret->children.emplace_back(child()->getStats()); + return ret; +} + +Status ReturnKeyStage::_extractIndexKey(WorkingSetMember* member) { + if (!_sortKeyMetaFields.empty()) { + invariant(member->metadata().hasSortKey()); + } + + auto indexKey = member->metadata().hasIndexKey() ? member->metadata().getIndexKey() : BSONObj(); + auto sortKey = member->metadata().hasSortKey() ? member->metadata().getSortKey() : BSONObj(); + BSONObjBuilder bob; + + if (!indexKey.isEmpty()) { + bob.appendElements(indexKey); + } + + for (const auto& fieldName : _sortKeyMetaFields) { + bob.append(fieldName, sortKey); + } + + member->keyData.clear(); + member->recordId = {}; + member->resetDocument({}, bob.obj()); + member->transitionToOwnedObj(); + + return Status::OK(); +} +} // namespace mongo diff --git a/src/mongo/db/exec/return_key.h b/src/mongo/db/exec/return_key.h new file mode 100644 index 00000000000..7212de12843 --- /dev/null +++ b/src/mongo/db/exec/return_key.h @@ -0,0 +1,83 @@ +/** + * 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" + +namespace mongo { +/** + * This stage returns the index key or keys, used for executing the query, for each document in + * the result. If the query does not use an index to perform the read operation, the stage returns + * empty documents. + * + * If 'sortKeyMetaFields' vector is specified while constructing this stage, each element in this + * array will be treated as a field name which will be added to the output document and will + * hold a sort key for the document in the result set, if the sort key exists. The elements in this + * array are the values specified in the 'sortKey' meta-projection. + */ +class ReturnKeyStage : public PlanStage { +public: + static constexpr StringData kStageName = "RETURN_KEY"_sd; + + ReturnKeyStage(OperationContext* opCtx, + std::vector<std::string> sortKeyMetaFields, + WorkingSet* ws, + std::unique_ptr<PlanStage> child) + : PlanStage(opCtx, std::move(child), kStageName.rawData()), + _ws(*ws), + _sortKeyMetaFields(std::move(sortKeyMetaFields)) {} + + StageType stageType() const final { + return STAGE_RETURN_KEY; + } + + bool isEOF() final { + return child()->isEOF(); + } + + StageState doWork(WorkingSetID* out) final; + + std::unique_ptr<PlanStageStats> getStats() final; + + const SpecificStats* getSpecificStats() const final { + return &_specificStats; + } + +private: + Status _extractIndexKey(WorkingSetMember* member); + + WorkingSet& _ws; + ReturnKeyStats _specificStats; + + // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey + // meta-projection. + std::vector<std::string> _sortKeyMetaFields; +}; +} // namespace mongo diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index 4b039fa7aaf..4d809846186 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -49,6 +49,7 @@ #include "mongo/db/exec/multi_plan.h" #include "mongo/db/exec/projection.h" #include "mongo/db/exec/record_store_fast_count.h" +#include "mongo/db/exec/return_key.h" #include "mongo/db/exec/shard_filter.h" #include "mongo/db/exec/sort_key_generator.h" #include "mongo/db/exec/subplan.h" @@ -388,23 +389,32 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx, std::move(root)); } - // 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. - if (nullptr != canonicalQuery->getProj()) { - - // Add a SortKeyGeneratorStage if there is a $meta sortKey projection. - if (canonicalQuery->getProj()->wantSortKey()) { - root = std::make_unique<SortKeyGeneratorStage>( - canonicalQuery->getExpCtx(), - std::move(root), - ws, - canonicalQuery->getQueryRequest().getSort()); - } + // Add a SortKeyGeneratorStage if there is a $meta sortKey projection. + if (canonicalQuery->getProj() && canonicalQuery->getProj()->wantSortKey()) { + root = std::make_unique<SortKeyGeneratorStage>( + canonicalQuery->getExpCtx(), + std::move(root), + ws, + canonicalQuery->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 + // returnKey. + root = std::make_unique<ReturnKeyStage>( + opCtx, + QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection( + canonicalQuery->getQueryRequest().getProj()), + ws, + std::move(root)); + } else if (canonicalQuery->getProj()) { + // 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 (canonicalQuery->getProj()->requiresDocument() || - canonicalQuery->getProj()->wantIndexKey() || canonicalQuery->getProj()->wantSortKey() || canonicalQuery->getProj()->hasDottedFieldPath()) { root = std::make_unique<ProjectionStageDefault>( diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp index 359ad5c23d8..165a27e66bb 100644 --- a/src/mongo/db/query/parsed_projection.cpp +++ b/src/mongo/db/query/parsed_projection.cpp @@ -56,8 +56,6 @@ Status ParsedProjection::make(OperationContext* opCtx, IncludeExclude includeExclude = IncludeExclude::kUninitialized; bool requiresDocument = false; - bool hasIndexKeyProjection = false; - bool wantTextScore = false; bool wantGeoNearPoint = false; bool wantGeoNearDistance = false; @@ -161,7 +159,6 @@ Status ParsedProjection::make(OperationContext* opCtx, if (e2.valuestr() != QueryRequest::metaTextScore && e2.valuestr() != QueryRequest::metaRecordId && - e2.valuestr() != QueryRequest::metaIndexKey && e2.valuestr() != QueryRequest::metaGeoNearDistance && e2.valuestr() != QueryRequest::metaGeoNearPoint && e2.valuestr() != QueryRequest::metaSortKey) { @@ -171,8 +168,6 @@ Status ParsedProjection::make(OperationContext* opCtx, // This clobbers everything else. if (e2.valuestr() == QueryRequest::metaTextScore) { wantTextScore = true; - } else if (e2.valuestr() == QueryRequest::metaIndexKey) { - hasIndexKeyProjection = true; } else if (e2.valuestr() == QueryRequest::metaGeoNearDistance) { wantGeoNearDistance = true; } else if (e2.valuestr() == QueryRequest::metaGeoNearPoint) { @@ -268,7 +263,6 @@ Status ParsedProjection::make(OperationContext* opCtx, // Save the raw spec. It should be owned by the QueryRequest. verify(spec.isOwned()); pp->_source = spec; - pp->_returnKey = hasIndexKeyProjection; pp->_requiresDocument = requiresDocument; // Add meta-projections. @@ -308,11 +302,6 @@ Status ParsedProjection::make(OperationContext* opCtx, } } - // returnKey clobbers everything except for sortKey meta-projection. - if (hasIndexKeyProjection && !wantSortKey) { - pp->_requiresDocument = false; - } - *out = pp.release(); return Status::OK(); } diff --git a/src/mongo/db/query/parsed_projection.h b/src/mongo/db/query/parsed_projection.h index 0a5ee6de3c2..1d9063acf05 100644 --- a/src/mongo/db/query/parsed_projection.h +++ b/src/mongo/db/query/parsed_projection.h @@ -100,10 +100,6 @@ public: return _wantGeoNearPoint; } - bool wantIndexKey() const { - return _returnKey; - } - bool wantSortKey() const { return _wantSortKey; } @@ -193,8 +189,6 @@ private: bool _wantGeoNearPoint = false; - bool _returnKey = false; - // Whether this projection includes a sortKey meta-projection. bool _wantSortKey = false; diff --git a/src/mongo/db/query/parsed_projection_test.cpp b/src/mongo/db/query/parsed_projection_test.cpp index 990b665d6ed..8994723a62d 100644 --- a/src/mongo/db/query/parsed_projection_test.cpp +++ b/src/mongo/db/query/parsed_projection_test.cpp @@ -234,7 +234,6 @@ TEST(ParsedProjectionTest, ParsedProjectionDefaults) { ASSERT_FALSE(parsedProjection->requiresMatchDetails()); ASSERT_FALSE(parsedProjection->wantGeoNearDistance()); ASSERT_FALSE(parsedProjection->wantGeoNearPoint()); - ASSERT_FALSE(parsedProjection->wantIndexKey()); } TEST(ParsedProjectionTest, SortKeyMetaProjection) { @@ -247,7 +246,6 @@ TEST(ParsedProjectionTest, SortKeyMetaProjection) { ASSERT_FALSE(parsedProjection->requiresMatchDetails()); ASSERT_FALSE(parsedProjection->wantGeoNearDistance()); ASSERT_FALSE(parsedProjection->wantGeoNearPoint()); - ASSERT_FALSE(parsedProjection->wantIndexKey()); } TEST(ParsedProjectionTest, SortKeyMetaProjectionCovered) { @@ -261,7 +259,6 @@ TEST(ParsedProjectionTest, SortKeyMetaProjectionCovered) { ASSERT_FALSE(parsedProjection->requiresMatchDetails()); ASSERT_FALSE(parsedProjection->wantGeoNearDistance()); ASSERT_FALSE(parsedProjection->wantGeoNearPoint()); - ASSERT_FALSE(parsedProjection->wantIndexKey()); } TEST(ParsedProjectionTest, SortKeyMetaAndSlice) { @@ -276,7 +273,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndSlice) { ASSERT_FALSE(parsedProjection->requiresMatchDetails()); ASSERT_FALSE(parsedProjection->wantGeoNearDistance()); ASSERT_FALSE(parsedProjection->wantGeoNearPoint()); - ASSERT_FALSE(parsedProjection->wantIndexKey()); } TEST(ParsedProjectionTest, SortKeyMetaAndElemMatch) { @@ -291,7 +287,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndElemMatch) { ASSERT_FALSE(parsedProjection->requiresMatchDetails()); ASSERT_FALSE(parsedProjection->wantGeoNearDistance()); ASSERT_FALSE(parsedProjection->wantGeoNearPoint()); - ASSERT_FALSE(parsedProjection->wantIndexKey()); } TEST(ParsedProjectionTest, SortKeyMetaAndExclusion) { @@ -305,7 +300,6 @@ TEST(ParsedProjectionTest, SortKeyMetaAndExclusion) { ASSERT_FALSE(parsedProjection->requiresMatchDetails()); ASSERT_FALSE(parsedProjection->wantGeoNearDistance()); ASSERT_FALSE(parsedProjection->wantGeoNearPoint()); - ASSERT_FALSE(parsedProjection->wantIndexKey()); } // diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp index 3487e955675..0d09172ad37 100644 --- a/src/mongo/db/query/planner_analysis.cpp +++ b/src/mongo/db/query/planner_analysis.cpp @@ -306,8 +306,8 @@ auto isCoveredOrAlreadyFetched(const vector<StringData>& fields, * Checks all properties that exclude a projection from being simple. */ auto isSimpleProjection(const CanonicalQuery& query) { - return !query.getProj()->wantIndexKey() && !query.getProj()->wantSortKey() && - !query.getProj()->hasDottedFieldPath() && !query.getProj()->requiresDocument(); + return !query.getProj()->wantSortKey() && !query.getProj()->hasDottedFieldPath() && + !query.getProj()->requiresDocument(); } /** @@ -332,6 +332,21 @@ auto produceCoveredKeyObj(QuerySolutionNode* solnRoot) { } /** + * Adds a stage to generate the sort key metadata if there's no sort stage but we have a sortKey + * meta-projection. + */ +std::unique_ptr<QuerySolutionNode> addSortKeyGeneratorStageIfNeeded( + const CanonicalQuery& query, bool hasSortStage, std::unique_ptr<QuerySolutionNode> solnRoot) { + if (!hasSortStage && query.getProj() && query.getProj()->wantSortKey()) { + auto keyGenNode = std::make_unique<SortKeyGeneratorNode>(); + keyGenNode->sortSpec = query.getQueryRequest().getSort(); + keyGenNode->children.push_back(solnRoot.release()); + return keyGenNode; + } + return solnRoot; +} + +/** * When projection needs to be added to the solution tree, this function chooses between the default * implementation and one of the fast paths. */ @@ -340,24 +355,12 @@ std::unique_ptr<ProjectionNode> analyzeProjection(const CanonicalQuery& query, const bool hasSortStage) { const QueryRequest& qr = query.getQueryRequest(); - // If there's no sort stage but we have a sortKey meta-projection, we need to add a stage to - // generate the sort key metadata. - auto addSortKeyGeneratorStageIfNeeded = [&]() { - if (!hasSortStage && query.getProj()->wantSortKey()) { - auto keyGenNode = std::make_unique<SortKeyGeneratorNode>(); - keyGenNode->sortSpec = qr.getSort(); - keyGenNode->children.push_back(solnRoot.release()); - solnRoot = std::move(keyGenNode); - } - }; - LOG(5) << "PROJECTION: Current plan is:\n" << redact(solnRoot->toString()); // If the projection requires the entire document we add a fetch stage if not present. Otherwise - // we add a fetch stage if we are not covered and not returnKey. + // we add a fetch stage if we are not covered. if ((query.getProj()->requiresDocument() && !solnRoot->fetched()) || - (!isCoveredOrAlreadyFetched(query.getProj()->getRequiredFields(), *solnRoot) && - !query.getProj()->wantIndexKey())) { + (!isCoveredOrAlreadyFetched(query.getProj()->getRequiredFields(), *solnRoot))) { auto fetch = std::make_unique<FetchNode>(); fetch->children.push_back(solnRoot.release()); solnRoot = std::move(fetch); @@ -372,30 +375,33 @@ std::unique_ptr<ProjectionNode> analyzeProjection(const CanonicalQuery& query, if (isSimpleProjection(query)) { // If the projection is simple, but not covered, use 'ProjectionNodeSimple'. if (solnRoot->fetched()) { - addSortKeyGeneratorStageIfNeeded(); return std::make_unique<ProjectionNodeSimple>( - std::move(solnRoot), *query.root(), qr.getProj(), *query.getProj()); + addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)), + *query.root(), + qr.getProj(), + *query.getProj()); } else { // If we're here we're not fetched so we're covered. Let's see if we can get out of // using the default projType. If 'solnRoot' is an index scan we can use the faster // covered impl. BSONObj coveredKeyObj = produceCoveredKeyObj(solnRoot.get()); if (!coveredKeyObj.isEmpty()) { - addSortKeyGeneratorStageIfNeeded(); - return std::make_unique<ProjectionNodeCovered>(std::move(solnRoot), - *query.root(), - qr.getProj(), - *query.getProj(), - std::move(coveredKeyObj)); + return std::make_unique<ProjectionNodeCovered>( + addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)), + *query.root(), + qr.getProj(), + *query.getProj(), + std::move(coveredKeyObj)); } } } - addSortKeyGeneratorStageIfNeeded(); return std::make_unique<ProjectionNodeDefault>( - std::move(solnRoot), *query.root(), qr.getProj(), *query.getProj()); + addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)), + *query.root(), + qr.getProj(), + *query.getProj()); } - } // namespace // static @@ -797,7 +803,13 @@ std::unique_ptr<QuerySolution> QueryPlannerAnalysis::analyzeDataAccess( } // Project the results. - if (query.getProj()) { + if (qr.returnKey()) { + // We don't need a projection stage if returnKey was requested since the intended behavior + // is that the projection is ignored when returnKey is specified. + solnRoot = std::make_unique<ReturnKeyNode>( + addSortKeyGeneratorStageIfNeeded(query, hasSortStage, std::move(solnRoot)), + QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection(qr.getProj())); + } else if (query.getProj()) { solnRoot = analyzeProjection(query, std::move(solnRoot), hasSortStage); // If we don't have a covered project, and we're not allowed to put an uncovered one in, // bail out. diff --git a/src/mongo/db/query/query_planner_common.cpp b/src/mongo/db/query/query_planner_common.cpp index 35cb1ac0293..2b2117676a8 100644 --- a/src/mongo/db/query/query_planner_common.cpp +++ b/src/mongo/db/query/query_planner_common.cpp @@ -75,4 +75,32 @@ void QueryPlannerCommon::reverseScans(QuerySolutionNode* node) { } } +// TODO SERVER-42422: reimplement by walking a projection AST rather than raw bson. +std::vector<std::string> QueryPlannerCommon::extractSortKeyMetaFieldsFromProjection( + const BSONObj& proj) { + std::vector<std::string> sortKeyMetaFields; + + for (auto&& elem : proj) { + if (elem.type() == BSONType::Object) { + BSONObj obj = elem.embeddedObject(); + // The caller must have already validated the projection, so we expect + // to see only one element. + invariant(1 == obj.nFields()); + + BSONElement firstSubElem = obj.firstElement(); + if (firstSubElem.fieldNameStringData() == "$meta") { + invariant(firstSubElem.type() == BSONType::String); + if (firstSubElem.valueStringData() == QueryRequest::metaSortKey) { + invariant(std::find(sortKeyMetaFields.begin(), + sortKeyMetaFields.end(), + elem.fieldName()) == sortKeyMetaFields.end()); + sortKeyMetaFields.push_back(elem.fieldName()); + } + } + } + } + + return sortKeyMetaFields; +} + } // namespace mongo diff --git a/src/mongo/db/query/query_planner_common.h b/src/mongo/db/query/query_planner_common.h index c232291e0e3..76c70742a3c 100644 --- a/src/mongo/db/query/query_planner_common.h +++ b/src/mongo/db/query/query_planner_common.h @@ -82,6 +82,14 @@ public: * the scan direction and index bounds. */ static void reverseScans(QuerySolutionNode* node); + + /** + * Extracts all field names for the sortKey meta-projection and stores them in the returned + * array. Returns an empty array if there were no sortKey meta-projection specified in the + * given projection 'proj'. For example, given a projection {a:1, b: {$meta: "sortKey"}, + * c: {$meta: "sortKey"}}, the returned vector will contain two elements ["b", "c"]. + */ + static std::vector<std::string> extractSortKeyMetaFieldsFromProjection(const BSONObj& proj); }; } // namespace mongo diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp index b3c87b40ab8..fde6441559b 100644 --- a/src/mongo/db/query/query_request.cpp +++ b/src/mongo/db/query/query_request.cpp @@ -58,7 +58,6 @@ const char QueryRequest::queryOptionMaxTimeMS[] = "$maxTimeMS"; const string QueryRequest::metaGeoNearDistance("geoNearDistance"); const string QueryRequest::metaGeoNearPoint("geoNearPoint"); -const string QueryRequest::metaIndexKey("indexKey"); const string QueryRequest::metaRecordId("recordId"); const string QueryRequest::metaSortKey("sortKey"); const string QueryRequest::metaTextScore("textScore"); @@ -592,16 +591,6 @@ void QueryRequest::asFindCommandInternal(BSONObjBuilder* cmdBuilder) const { } } -void QueryRequest::addReturnKeyMetaProj() { - BSONObjBuilder projBob; - projBob.appendElements(_proj); - // We use $$ because it's never going to show up in a user's projection. - // The exact text doesn't matter. - BSONObj indexKey = BSON("$$" << BSON("$meta" << QueryRequest::metaIndexKey)); - projBob.append(indexKey.firstElement()); - _proj = projBob.obj(); -} - void QueryRequest::addShowRecordIdMetaProj() { BSONObjBuilder projBob; projBob.appendElements(_proj); @@ -937,7 +926,6 @@ Status QueryRequest::initFullQuery(const BSONObj& top) { // Won't throw. if (e.trueValue()) { _returnKey = true; - addReturnKeyMetaProj(); } } else if (name == "showDiskLoc") { // Won't throw. @@ -1004,11 +992,6 @@ void QueryRequest::initFromInt(int options) { } void QueryRequest::addMetaProjection() { - // We might need to update the projection object with a $meta projection. - if (returnKey()) { - addReturnKeyMetaProj(); - } - if (showRecordId()) { addShowRecordIdMetaProj(); } diff --git a/src/mongo/db/query/query_request.h b/src/mongo/db/query/query_request.h index 7b0cb74dd80..17e915ace5c 100644 --- a/src/mongo/db/query/query_request.h +++ b/src/mongo/db/query/query_request.h @@ -135,7 +135,6 @@ public: // Names of the $meta projection values. static const std::string metaGeoNearDistance; static const std::string metaGeoNearPoint; - static const std::string metaIndexKey; static const std::string metaRecordId; static const std::string metaSortKey; static const std::string metaTextScore; @@ -455,11 +454,6 @@ private: Status initFullQuery(const BSONObj& top); /** - * Updates the projection object with a $meta projection for the returnKey option. - */ - void addReturnKeyMetaProj(); - - /** * Updates the projection object with a $meta projection for the showRecordId option. */ void addShowRecordIdMetaProj(); diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp index ddbe87074a8..6b076d80e96 100644 --- a/src/mongo/db/query/query_solution.cpp +++ b/src/mongo/db/query/query_solution.cpp @@ -31,6 +31,8 @@ #include "mongo/db/query/query_solution.h" +#include <boost/algorithm/string/join.hpp> + #include "mongo/bson/bsontypes.h" #include "mongo/bson/mutable/document.h" #include "mongo/bson/simple_bsonelement_comparator.h" @@ -848,6 +850,27 @@ bool IndexScanNode::operator==(const IndexScanNode& other) const { } // +// ReturnKeyNode +// + +void ReturnKeyNode::appendToString(str::stream* ss, int indent) const { + addIndent(ss, indent); + *ss << "RETURN_KEY\n"; + addIndent(ss, indent + 1); + *ss << "sortKeyMetaFields = [" << boost::algorithm::join(sortKeyMetaFields, ", ") << "]\n"; + addCommon(ss, indent); + addIndent(ss, indent + 1); + *ss << "Child:" << '\n'; + children[0]->appendToString(ss, indent + 2); +} + +QuerySolutionNode* ReturnKeyNode::clone() const { + auto copy = std::make_unique<ReturnKeyNode>( + std::unique_ptr<QuerySolutionNode>(children[0]->clone()), std::vector(sortKeyMetaFields)); + return copy.release(); +} + +// // ProjectionNode // diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h index cba57fd1263..48a0288a317 100644 --- a/src/mongo/db/query/query_solution.h +++ b/src/mongo/db/query/query_solution.h @@ -540,6 +540,35 @@ struct IndexScanNode : public QuerySolutionNode { std::set<StringData> multikeyFields; }; +struct ReturnKeyNode : public QuerySolutionNode { + ReturnKeyNode(std::unique_ptr<QuerySolutionNode> child, + std::vector<std::string> sortKeyMetaFields) + : QuerySolutionNode(std::move(child)), sortKeyMetaFields(std::move(sortKeyMetaFields)) {} + + StageType getType() const final { + return STAGE_RETURN_KEY; + } + + void appendToString(str::stream* ss, int indent) const final; + + bool fetched() const final { + return children[0]->fetched(); + } + bool hasField(const std::string& field) const final { + return false; + } + bool sortedByDiskLoc() const final { + return children[0]->sortedByDiskLoc(); + } + const BSONObjSet& getSort() const final { + return children[0]->getSort(); + } + + QuerySolutionNode* clone() const final; + + std::vector<std::string> sortKeyMetaFields; +}; + /** * We have a few implementations of the projection functionality. They are chosen by constructing * a type derived from this abstract struct. The most general implementation 'ProjectionNodeDefault' diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp index 48cc9e39027..c8a92a75963 100644 --- a/src/mongo/db/query/stage_builder.cpp +++ b/src/mongo/db/query/stage_builder.cpp @@ -52,6 +52,7 @@ #include "mongo/db/exec/merge_sort.h" #include "mongo/db/exec/or.h" #include "mongo/db/exec/projection.h" +#include "mongo/db/exec/return_key.h" #include "mongo/db/exec/shard_filter.h" #include "mongo/db/exec/skip.h" #include "mongo/db/exec/sort.h" @@ -133,6 +134,13 @@ std::unique_ptr<PlanStage> buildStages(OperationContext* opCtx, return std::make_unique<SortKeyGeneratorStage>( 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); + return std::make_unique<ReturnKeyStage>( + opCtx, 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); diff --git a/src/mongo/db/query/stage_types.h b/src/mongo/db/query/stage_types.h index 8b34d4dfecc..17a67f1169f 100644 --- a/src/mongo/db/query/stage_types.h +++ b/src/mongo/db/query/stage_types.h @@ -87,6 +87,7 @@ enum StageType { STAGE_QUEUED_DATA, STAGE_RECORD_STORE_FAST_COUNT, + STAGE_RETURN_KEY, STAGE_SHARDING_FILTER, STAGE_SKIP, STAGE_SORT, |