SERVER-19355 add {$meta: 'sortKey'} projection

10 files changed, 118 insertions, 17 deletions

diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index 427c2083f2d..497d9151938 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -335,6 +335,10 @@ Status CanonicalQuery::init(LiteParsedQuery* lpq,
         _proj.reset(pp);
     }
 
+    if (_proj && _proj->wantSortKey() && _pq->getSort().isEmpty()) {
+        return Status(ErrorCodes::BadValue, "cannot use sortKey $meta projection without a sort");
+    }
+
     return Status::OK();
 }
 
diff --git a/src/mongo/db/query/canonical_query_test.cpp b/src/mongo/db/query/canonical_query_test.cpp
index 34a4575bde7..eff6201663d 100644
--- a/src/mongo/db/query/canonical_query_test.cpp
+++ b/src/mongo/db/query/canonical_query_test.cpp
@@ -399,6 +399,28 @@ TEST(CanonicalQueryTest, IsValidTextAndSnapshot) {
     ASSERT_NOT_OK(isValid("{$text: {$search: 's'}}", *lpq));
 }
 
+TEST(CanonicalQueryTest, IsValidSortKeyMetaProjection) {
+    // Passing a sortKey meta-projection without a sort is an error.
+    {
+        const bool isExplain = false;
+        auto lpq = assertGet(LiteParsedQuery::makeFromFindCommand(
+            nss, fromjson("{find: 'testcoll', projection: {foo: {$meta: 'sortKey'}}}"), isExplain));
+        auto cq = CanonicalQuery::canonicalize(lpq.release());
+        ASSERT_NOT_OK(cq.getStatus());
+    }
+
+    // Should be able to successfully create a CQ when there is a sort.
+    {
+        const bool isExplain = false;
+        auto lpq = assertGet(LiteParsedQuery::makeFromFindCommand(
+            nss,
+            fromjson("{find: 'testcoll', projection: {foo: {$meta: 'sortKey'}}, sort: {bar: 1}}"),
+            isExplain));
+        auto cq = CanonicalQuery::canonicalize(lpq.release());
+        ASSERT_OK(cq.getStatus());
+    }
+}
+
 //
 // Tests for CanonicalQuery::sortTree
 //
diff --git a/src/mongo/db/query/lite_parsed_query.cpp b/src/mongo/db/query/lite_parsed_query.cpp
index de93f29067f..96fe9a392e9 100644
--- a/src/mongo/db/query/lite_parsed_query.cpp
+++ b/src/mongo/db/query/lite_parsed_query.cpp
@@ -49,11 +49,12 @@ const char* LiteParsedQuery::kFindCommandReadPrefField("$readPreference");
 const string LiteParsedQuery::cmdOptionMaxTimeMS("maxTimeMS");
 const string LiteParsedQuery::queryOptionMaxTimeMS("$maxTimeMS");
 
-const string LiteParsedQuery::metaTextScore("textScore");
 const string LiteParsedQuery::metaGeoNearDistance("geoNearDistance");
 const string LiteParsedQuery::metaGeoNearPoint("geoNearPoint");
-const string LiteParsedQuery::metaRecordId("recordId");
 const string LiteParsedQuery::metaIndexKey("indexKey");
+const string LiteParsedQuery::metaRecordId("recordId");
+const string LiteParsedQuery::metaSortKey("sortKey");
+const string LiteParsedQuery::metaTextScore("textScore");
 
 const long long LiteParsedQuery::kDefaultBatchSize = 101;
 
diff --git a/src/mongo/db/query/lite_parsed_query.h b/src/mongo/db/query/lite_parsed_query.h
index ab1f8ab3aeb..76f10438101 100644
--- a/src/mongo/db/query/lite_parsed_query.h
+++ b/src/mongo/db/query/lite_parsed_query.h
@@ -158,11 +158,12 @@ public:
     static const std::string queryOptionMaxTimeMS;
 
     // Names of the $meta projection values.
-    static const std::string metaTextScore;
     static const std::string metaGeoNearDistance;
     static const std::string metaGeoNearPoint;
-    static const std::string metaRecordId;
     static const std::string metaIndexKey;
+    static const std::string metaRecordId;
+    static const std::string metaSortKey;
+    static const std::string metaTextScore;
 
     const NamespaceString& nss() const {
         return _nss;
diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp
index 6c01b0b3434..32979e9693e 100644
--- a/src/mongo/db/query/parsed_projection.cpp
+++ b/src/mongo/db/query/parsed_projection.cpp
@@ -66,6 +66,7 @@ Status ParsedProjection::make(const BSONObj& spec,
 
     bool wantGeoNearPoint = false;
     bool wantGeoNearDistance = false;
+    bool wantSortKey = false;
 
     // Until we see a positional or elemMatch operator we're normal.
     ArrayOpType arrayOpType = ARRAY_OP_NORMAL;
@@ -150,7 +151,8 @@ Status ParsedProjection::make(const BSONObj& spec,
                     e2.valuestr() != LiteParsedQuery::metaRecordId &&
                     e2.valuestr() != LiteParsedQuery::metaIndexKey &&
                     e2.valuestr() != LiteParsedQuery::metaGeoNearDistance &&
-                    e2.valuestr() != LiteParsedQuery::metaGeoNearPoint) {
+                    e2.valuestr() != LiteParsedQuery::metaGeoNearPoint &&
+                    e2.valuestr() != LiteParsedQuery::metaSortKey) {
                     return Status(ErrorCodes::BadValue, "unsupported $meta operator: " + e2.str());
                 }
 
@@ -161,6 +163,8 @@ Status ParsedProjection::make(const BSONObj& spec,
                     wantGeoNearDistance = true;
                 } else if (e2.valuestr() == LiteParsedQuery::metaGeoNearPoint) {
                     wantGeoNearPoint = true;
+                } else if (e2.valuestr() == LiteParsedQuery::metaSortKey) {
+                    wantSortKey = true;
                 }
             } else {
                 return Status(ErrorCodes::BadValue,
@@ -242,9 +246,10 @@ Status ParsedProjection::make(const BSONObj& spec,
     // missing."
     pp->_requiresDocument = include || hasNonSimple || hasDottedField;
 
-    // Add geoNear projections.
+    // Add meta-projections.
     pp->_wantGeoNearPoint = wantGeoNearPoint;
     pp->_wantGeoNearDistance = wantGeoNearDistance;
+    pp->_wantSortKey = wantSortKey;
 
     // If it's possible to compute the projection in a covered fashion, populate _requiredFields
     // so the planner can perform projection analysis.
diff --git a/src/mongo/db/query/parsed_projection.h b/src/mongo/db/query/parsed_projection.h
index b135b5f47ed..8c76944b796 100644
--- a/src/mongo/db/query/parsed_projection.h
+++ b/src/mongo/db/query/parsed_projection.h
@@ -96,16 +96,15 @@ public:
         return _returnKey;
     }
 
+    bool wantSortKey() const {
+        return _wantSortKey;
+    }
+
 private:
     /**
      * Must go through ::make
      */
-    ParsedProjection()
-        : _requiresMatchDetails(false),
-          _requiresDocument(true),
-          _wantGeoNearDistance(false),
-          _wantGeoNearPoint(false),
-          _returnKey(false) {}
+    ParsedProjection() = default;
 
     /**
      * Returns true if field name refers to a positional projection.
@@ -127,17 +126,20 @@ private:
     // TODO: stringdata?
     std::vector<std::string> _requiredFields;
 
-    bool _requiresMatchDetails;
+    bool _requiresMatchDetails = false;
 
-    bool _requiresDocument;
+    bool _requiresDocument = true;
 
     BSONObj _source;
 
-    bool _wantGeoNearDistance;
+    bool _wantGeoNearDistance = false;
+
+    bool _wantGeoNearPoint = false;
 
-    bool _wantGeoNearPoint;
+    bool _returnKey = false;
 
-    bool _returnKey;
+    // Whether this projection includes a sortKey meta-projection.
+    bool _wantSortKey = false;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/parsed_projection_test.cpp b/src/mongo/db/query/parsed_projection_test.cpp
index 4db3ebebe48..f850221b5fe 100644
--- a/src/mongo/db/query/parsed_projection_test.cpp
+++ b/src/mongo/db/query/parsed_projection_test.cpp
@@ -191,6 +191,30 @@ TEST(ParsedProjectionTest, InvalidPositionalProjectionDefaultPathMatchExpression
     ASSERT(!status.isOK());
 }
 
+TEST(ParsedProjectionTest, ParsedProjectionDefaults) {
+    auto parsedProjection = createParsedProjection("{}", "{}");
+
+    ASSERT_FALSE(parsedProjection->wantSortKey());
+    ASSERT_TRUE(parsedProjection->requiresDocument());
+    ASSERT_FALSE(parsedProjection->requiresMatchDetails());
+    ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
+    ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
+    ASSERT_FALSE(parsedProjection->wantIndexKey());
+}
+
+TEST(ParsedProjectionTest, SortKeyMetaProjection) {
+    auto parsedProjection = createParsedProjection("{}", "{foo: {$meta: 'sortKey'}}");
+
+    ASSERT_EQ(parsedProjection->getProjObj(), fromjson("{foo: {$meta: 'sortKey'}}"));
+    ASSERT_TRUE(parsedProjection->wantSortKey());
+    ASSERT_TRUE(parsedProjection->requiresDocument());
+
+    ASSERT_FALSE(parsedProjection->requiresMatchDetails());
+    ASSERT_FALSE(parsedProjection->wantGeoNearDistance());
+    ASSERT_FALSE(parsedProjection->wantGeoNearPoint());
+    ASSERT_FALSE(parsedProjection->wantIndexKey());
+}
+
 //
 // DBRef projections
 //
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 2142b1b781d..cf079fd140b 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -762,6 +762,16 @@ QuerySolution* QueryPlannerAnalysis::analyzeDataAccess(const CanonicalQuery& que
             }
         }
 
+        // If there's no sort stage but we have a sortKey meta-projection, we need to add a stage to
+        // generate the sort key computed data.
+        if (!hasSortStage && query.getProj()->wantSortKey()) {
+            SortKeyGeneratorNode* keyGenNode = new SortKeyGeneratorNode();
+            keyGenNode->queryObj = lpq.getFilter();
+            keyGenNode->sortSpec = lpq.getSort();
+            keyGenNode->children.push_back(solnRoot);
+            solnRoot = keyGenNode;
+        }
+
         // We now know we have whatever data is required for the projection.
         ProjectionNode* projNode = new ProjectionNode();
         projNode->children.push_back(solnRoot);
diff --git a/src/mongo/db/query/query_planner_test.cpp b/src/mongo/db/query/query_planner_test.cpp
index b7c68bba37c..3d3a07e9e45 100644
--- a/src/mongo/db/query/query_planner_test.cpp
+++ b/src/mongo/db/query/query_planner_test.cpp
@@ -3921,6 +3921,22 @@ TEST_F(QueryPlannerTest, CoveredOrUniqueIndexLookup) {
         "{ixscan: {filter: null, pattern: {a: 1, b: 1}}}}}");
 }
 
+TEST_F(QueryPlannerTest, SortKeyMetaProjection) {
+    addIndex(BSON("a" << 1));
+
+    runQuerySortProj(BSONObj(), fromjson("{a: 1}"), fromjson("{b: {$meta: 'sortKey'}}"));
+
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{proj: {spec: {b: {$meta: 'sortKey'}}, node: "
+        "{sort: {limit: 0, pattern: {a: 1}, node: {sortKeyGen: {node: "
+        "{cscan: {dir: 1}}}}}}}}");
+    assertSolutionExists(
+        "{proj: {spec: {b: {$meta: 'sortKey'}}, node: "
+        "{sortKeyGen: {node: {fetch: {filter: null, node: "
+        "{ixscan: {pattern: {a: 1}}}}}}}}}");
+}
+
 //
 // Test bad input to query planner helpers.
 //
diff --git a/src/mongo/db/query/query_planner_text_test.cpp b/src/mongo/db/query/query_planner_text_test.cpp
index 2a041e7008f..884b8d1dac9 100644
--- a/src/mongo/db/query/query_planner_text_test.cpp
+++ b/src/mongo/db/query/query_planner_text_test.cpp
@@ -412,4 +412,20 @@ TEST_F(QueryPlannerTest, TextDiacriticSensitive) {
     assertSolutionExists("{text: {search: 'blah', diacriticSensitive: true}}");
 }
 
+TEST_F(QueryPlannerTest, SortKeyMetaProjectionWithTextScoreMetaSort) {
+    addIndex(BSON("_fts"
+                  << "text"
+                  << "_ftsx" << 1));
+
+    runQuerySortProj(fromjson("{$text: {$search: 'foo'}}"),
+                     fromjson("{a: {$meta: 'textScore'}}"),
+                     fromjson("{a: {$meta: 'textScore'}, b: {$meta: 'sortKey'}}"));
+
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{proj: {spec: {a: {$meta: 'textScore'}, b: {$meta: 'sortKey'}}, node: "
+        "{sort: {limit: 0, pattern: {a: {$meta: 'textScore'}}, node: "
+        "{sortKeyGen: {node: {text: {search: 'foo'}}}}}}}}");
+}
+
 }  // namespace