summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--jstests/core/return_key.js76
-rw-r--r--src/mongo/db/exec/projection_exec.cpp53
-rw-r--r--src/mongo/db/exec/projection_exec.h8
-rw-r--r--src/mongo/db/query/parsed_projection.cpp4
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;
}