diff options
-rw-r--r-- | jstests/core/return_key.js | 76 | ||||
-rw-r--r-- | src/mongo/db/exec/projection_exec.cpp | 53 | ||||
-rw-r--r-- | src/mongo/db/exec/projection_exec.h | 8 | ||||
-rw-r--r-- | src/mongo/db/query/parsed_projection.cpp | 4 |
4 files changed, 124 insertions, 17 deletions
diff --git a/jstests/core/return_key.js b/jstests/core/return_key.js new file mode 100644 index 00000000000..cf02357759d --- /dev/null +++ b/jstests/core/return_key.js @@ -0,0 +1,76 @@ +/** + * Tests for returnKey. + */ +load("jstests/libs/analyze_plan.js"); + +(function() { + 'use strict'; + + var results; + var explain; + + var coll = db.jstests_returnkey; + coll.drop(); + + assert.writeOK(coll.insert({a: 1, b: 3})); + assert.writeOK(coll.insert({a: 2, b: 2})); + assert.writeOK(coll.insert({a: 3, b: 1})); + + assert.commandWorked(coll.ensureIndex({a: 1})); + assert.commandWorked(coll.ensureIndex({b: 1})); + + // Basic returnKey. + results = coll.find().hint({a: 1}).sort({a: 1}).returnKey().toArray(); + assert.eq(results, [{a: 1}, {a: 2}, {a: 3}]); + results = coll.find().hint({a: 1}).sort({a: -1}).returnKey().toArray(); + assert.eq(results, [{a: 3}, {a: 2}, {a: 1}]); + + // Check that the plan is covered. + explain = coll.find().hint({a: 1}).sort({a: 1}).returnKey().explain(); + assert(isIndexOnly(explain.queryPlanner.winningPlan)); + explain = coll.find().hint({a: 1}).sort({a: -1}).returnKey().explain(); + assert(isIndexOnly(explain.queryPlanner.winningPlan)); + + // returnKey with an in-memory sort. + results = coll.find().hint({a: 1}).sort({b: 1}).returnKey().toArray(); + assert.eq(results, [{a: 3}, {a: 2}, {a: 1}]); + results = coll.find().hint({a: 1}).sort({b: -1}).returnKey().toArray(); + assert.eq(results, [{a: 1}, {a: 2}, {a: 3}]); + + // Check that the plan is not covered. + explain = coll.find().hint({a: 1}).sort({b: 1}).returnKey().explain(); + assert(!isIndexOnly(explain.queryPlanner.winningPlan)); + explain = coll.find().hint({a: 1}).sort({b: -1}).returnKey().explain(); + assert(!isIndexOnly(explain.queryPlanner.winningPlan)); + + // returnKey takes precedence over other a regular inclusion projection. Should still be + // covered. + results = coll.find({}, {b: 1}).hint({a: 1}).sort({a: -1}).returnKey().toArray(); + assert.eq(results, [{a: 3}, {a: 2}, {a: 1}]); + explain = coll.find({}, {b: 1}).hint({a: 1}).sort({a: -1}).returnKey().explain(); + assert(isIndexOnly(explain.queryPlanner.winningPlan)); + + // returnKey takes precedence over other a regular exclusion projection. Should still be + // covered. + results = coll.find({}, {a: 0}).hint({a: 1}).sort({a: -1}).returnKey().toArray(); + assert.eq(results, [{a: 3}, {a: 2}, {a: 1}]); + explain = coll.find({}, {a: 0}).hint({a: 1}).sort({a: -1}).returnKey().explain(); + assert(isIndexOnly(explain.queryPlanner.winningPlan)); + + // Unlike other projections, sortKey meta-projection can co-exist with returnKey. + results = coll.find({}, {c: {$meta: 'sortKey'}}) + .hint({a: 1}).sort({a: -1}).returnKey().toArray(); + assert.eq(results, [{a: 3, c: {'': 3}}, {a: 2, c: {'': 2}}, {a: 1, c: {'': 1}}]); + + // returnKey with sortKey $meta where there is an in-memory sort. + results = coll.find({}, {c: {$meta: 'sortKey'}}) + .hint({a: 1}).sort({b: 1}).returnKey().toArray(); + assert.eq(results, [{a: 3, c: {'': 1}}, {a: 2, c: {'': 2}}, {a: 1, c: {'': 3}}]); + + // returnKey with multiple sortKey $meta projections. + results = coll.find({}, {c: {$meta: 'sortKey'}, d: {$meta: 'sortKey'}}) + .hint({a: 1}).sort({b: 1}).returnKey().toArray(); + assert.eq(results, [{a: 3, c: {'': 1}, d: {'': 1}}, + {a: 2, c: {'': 2}, d: {'': 2}}, + {a: 1, c: {'': 3}, d: {'': 3}}]); +})(); diff --git a/src/mongo/db/exec/projection_exec.cpp b/src/mongo/db/exec/projection_exec.cpp index fc20f4e53bf..1d996d1c6eb 100644 --- a/src/mongo/db/exec/projection_exec.cpp +++ b/src/mongo/db/exec/projection_exec.cpp @@ -39,6 +39,29 @@ namespace mongo { using std::max; using std::string; +namespace { + +/** + * Adds sort key metadata inside 'member' to 'builder' with field name 'fieldName'. + * + * Returns a non-OK status if sort key metadata is missing from 'member'. + */ +Status addSortKeyMetaProj(StringData fieldName, + const WorkingSetMember& member, + BSONObjBuilder* builder) { + if (!member.hasComputed(WSM_SORT_KEY)) { + return Status(ErrorCodes::InternalError, + "sortKey meta-projection requested but no data available"); + } + + const SortKeyComputedData* sortKeyData = + static_cast<const SortKeyComputedData*>(member.getComputed(WSM_SORT_KEY)); + builder->append(fieldName, sortKeyData->getSortKey()); + return Status::OK(); +} + +} // namespace + ProjectionExec::ProjectionExec() : _include(true), _special(false), @@ -51,7 +74,6 @@ ProjectionExec::ProjectionExec() _queryExpression(NULL), _hasReturnKey(false) {} - ProjectionExec::ProjectionExec(const BSONObj& spec, const MatchExpression* queryExpression, const MatchExpressionParser::WhereCallback& whereCallback) @@ -126,7 +148,8 @@ ProjectionExec::ProjectionExec(const BSONObj& spec, if (e2.valuestr() == LiteParsedQuery::metaTextScore) { _meta[e.fieldName()] = META_TEXT_SCORE; } else if (e2.valuestr() == LiteParsedQuery::metaSortKey) { - _meta[e.fieldName()] = META_SORT_KEY; + _sortKeyMetaFields.push_back(e.fieldName()); + _meta[_sortKeyMetaFields.back()] = META_SORT_KEY; } else if (e2.valuestr() == LiteParsedQuery::metaRecordId) { _meta[e.fieldName()] = META_RECORDID; } else if (e2.valuestr() == LiteParsedQuery::metaGeoNearPoint) { @@ -135,8 +158,6 @@ ProjectionExec::ProjectionExec(const BSONObj& spec, _meta[e.fieldName()] = META_GEONEAR_DIST; } else if (e2.valuestr() == LiteParsedQuery::metaIndexKey) { _hasReturnKey = true; - // The index key clobbers everything so just stop parsing here. - return; } else { // This shouldn't happen, should be caught by parsing. verify(0); @@ -226,15 +247,24 @@ void ProjectionExec::add(const string& field, int skip, int limit) { Status ProjectionExec::transform(WorkingSetMember* member) const { if (_hasReturnKey) { - BSONObj keyObj; + BSONObjBuilder builder; if (member->hasComputed(WSM_INDEX_KEY)) { const IndexKeyComputedData* key = static_cast<const IndexKeyComputedData*>(member->getComputed(WSM_INDEX_KEY)); - keyObj = key->getKey(); + builder.appendElements(key->getKey()); } - member->obj = Snapshotted<BSONObj>(SnapshotId(), keyObj.getOwned()); + // Must be possible to do both returnKey meta-projection and sortKey meta-projection so that + // mongos can support returnKey. + for (auto fieldName : _sortKeyMetaFields) { + auto sortKeyMetaStatus = addSortKeyMetaProj(fieldName, *member, &builder); + if (!sortKeyMetaStatus.isOK()) { + return sortKeyMetaStatus; + } + } + + member->obj = Snapshotted<BSONObj>(SnapshotId(), builder.obj()); member->keyData.clear(); member->loc = RecordId(); member->transitionToOwnedObj(); @@ -316,13 +346,10 @@ Status ProjectionExec::transform(WorkingSetMember* member) const { bob.append(it->first, 0.0); } } else if (META_SORT_KEY == it->second) { - if (!member->hasComputed(WSM_SORT_KEY)) { - return Status(ErrorCodes::InternalError, - "sortKey meta-projection requested but no data available"); + auto sortKeyMetaStatus = addSortKeyMetaProj(it->first, *member, &bob); + if (!sortKeyMetaStatus.isOK()) { + return sortKeyMetaStatus; } - const SortKeyComputedData* sortKeyData = - static_cast<const SortKeyComputedData*>(member->getComputed(WSM_SORT_KEY)); - bob.append(it->first, sortKeyData->getSortKey()); } else if (META_RECORDID == it->second) { bob.append(it->first, static_cast<long long>(member->loc.repr())); } diff --git a/src/mongo/db/exec/projection_exec.h b/src/mongo/db/exec/projection_exec.h index 8079ee8afa8..906d2fe12d8 100644 --- a/src/mongo/db/exec/projection_exec.h +++ b/src/mongo/db/exec/projection_exec.h @@ -190,9 +190,13 @@ 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. If - // it's not found we output nothing. + // 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; + + // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey + // meta-projection. + std::vector<StringData> _sortKeyMetaFields; }; } // namespace mongo diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp index 32979e9693e..4f6a48c2f9e 100644 --- a/src/mongo/db/query/parsed_projection.cpp +++ b/src/mongo/db/query/parsed_projection.cpp @@ -273,8 +273,8 @@ Status ParsedProjection::make(const BSONObj& spec, } } - // returnKey clobbers everything. - if (hasIndexKeyProjection) { + // returnKey clobbers everything except for sortKey meta-projection. + if (hasIndexKeyProjection && !wantSortKey) { pp->_requiresDocument = false; } |