SERVER-42298 Use SortPattern in the implementation of SortKeyGenerator

18 files changed, 166 insertions, 139 deletions

diff --git a/src/mongo/db/exec/sort_key_generator.cpp b/src/mongo/db/exec/sort_key_generator.cpp
index 7b846f18bad..90fe3b48fea 100644
--- a/src/mongo/db/exec/sort_key_generator.cpp
+++ b/src/mongo/db/exec/sort_key_generator.cpp
@@ -49,12 +49,13 @@ namespace mongo {
 
 const char* SortKeyGeneratorStage::kStageType = "SORT_KEY_GENERATOR";
 
-SortKeyGeneratorStage::SortKeyGeneratorStage(OperationContext* opCtx,
+SortKeyGeneratorStage::SortKeyGeneratorStage(const boost::intrusive_ptr<ExpressionContext>& pExpCtx,
                                              PlanStage* child,
                                              WorkingSet* ws,
-                                             const BSONObj& sortSpecObj,
-                                             const CollatorInterface* collator)
-    : PlanStage(kStageType, opCtx), _ws(ws), _sortKeyGen(sortSpecObj, collator) {
+                                             const BSONObj& sortSpecObj)
+    : PlanStage(kStageType, pExpCtx->opCtx),
+      _ws(ws),
+      _sortKeyGen({{sortSpecObj, pExpCtx}, pExpCtx->getCollator()}) {
     _children.emplace_back(child);
 }
 
diff --git a/src/mongo/db/exec/sort_key_generator.h b/src/mongo/db/exec/sort_key_generator.h
index 79cb5a6fe49..c7a9e2dfd4b 100644
--- a/src/mongo/db/exec/sort_key_generator.h
+++ b/src/mongo/db/exec/sort_key_generator.h
@@ -34,6 +34,7 @@
 #include "mongo/bson/bsonobj.h"
 #include "mongo/db/exec/plan_stage.h"
 #include "mongo/db/index/sort_key_generator.h"
+#include "mongo/db/pipeline/expression_context.h"
 #include "mongo/db/query/index_bounds.h"
 #include "mongo/db/query/stage_types.h"
 
@@ -49,11 +50,10 @@ class WorkingSetMember;
  */
 class SortKeyGeneratorStage final : public PlanStage {
 public:
-    SortKeyGeneratorStage(OperationContext* opCtx,
+    SortKeyGeneratorStage(const boost::intrusive_ptr<ExpressionContext>& pExpCtx,
                           PlanStage* child,
                           WorkingSet* ws,
-                          const BSONObj& sortSpecObj,
-                          const CollatorInterface* collator);
+                          const BSONObj& sortSpecObj);
 
     bool isEOF() final;
 
diff --git a/src/mongo/db/exec/sort_test.cpp b/src/mongo/db/exec/sort_test.cpp
index 66d6f3c1d86..7c8a1696399 100644
--- a/src/mongo/db/exec/sort_test.cpp
+++ b/src/mongo/db/exec/sort_test.cpp
@@ -103,8 +103,12 @@ public:
         params.pattern = fromjson(patternStr);
         params.limit = limit;
 
+        // Create an ExpressionContext for the SortKeyGeneratorStage.
+        boost::intrusive_ptr<ExpressionContext> pExpCtx(
+            new ExpressionContext(getOpCtx(), collator));
+
         auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-            getOpCtx(), queuedDataStage.release(), &ws, params.pattern, collator);
+            pExpCtx, queuedDataStage.release(), &ws, params.pattern);
 
         SortStage sort(getOpCtx(), params, &ws, sortKeyGen.release());
 
@@ -158,10 +162,13 @@ private:
 TEST_F(SortStageTest, SortEmptyWorkingSet) {
     WorkingSet ws;
 
+    // Create an ExpressionContext for the SortKeyGeneratorStage.
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(getOpCtx(), nullptr));
+
     // QueuedDataStage will be owned by SortStage.
     auto queuedDataStage = std::make_unique<QueuedDataStage>(getOpCtx(), &ws);
-    auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-        getOpCtx(), queuedDataStage.release(), &ws, BSONObj(), nullptr);
+    auto sortKeyGen =
+        std::make_unique<SortKeyGeneratorStage>(pExpCtx, queuedDataStage.release(), &ws, BSONObj());
     SortStageParams params;
     SortStage sort(getOpCtx(), params, &ws, sortKeyGen.release());
 
diff --git a/src/mongo/db/index/SConscript b/src/mongo/db/index/SConscript
index 07784425da7..3d7f912bd6f 100644
--- a/src/mongo/db/index/SConscript
+++ b/src/mongo/db/index/SConscript
@@ -53,6 +53,7 @@ env.Library(
             '$BUILD_DIR/mongo/db/mongohasher',
             '$BUILD_DIR/mongo/db/projection_exec_agg',
             '$BUILD_DIR/mongo/db/query/collation/collator_interface',
+            '$BUILD_DIR/mongo/db/query/sort_pattern',
             '$BUILD_DIR/third_party/s2/s2',
             'expression_params',
             'index_descriptor',
@@ -163,5 +164,6 @@ env.CppUnitTest(
         "$BUILD_DIR/mongo/db/matcher/expressions",
         '$BUILD_DIR/mongo/db/mongohasher',
         '$BUILD_DIR/mongo/db/query/collation/collator_interface_mock',
+        '$BUILD_DIR/mongo/db/query/query_test_service_context',
     ],
 )
diff --git a/src/mongo/db/index/sort_key_generator.cpp b/src/mongo/db/index/sort_key_generator.cpp
index d07f4d9bc62..82122a65448 100644
--- a/src/mongo/db/index/sort_key_generator.cpp
+++ b/src/mongo/db/index/sort_key_generator.cpp
@@ -36,40 +36,23 @@
 
 namespace mongo {
 
-SortKeyGenerator::SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterface* collator)
-    : _collator(collator) {
+SortKeyGenerator::SortKeyGenerator(SortPattern sortPattern, const CollatorInterface* collator)
+    : _collator(collator), _sortPattern(std::move(sortPattern)) {
     BSONObjBuilder btreeBob;
+    size_t nFields = 0;
 
-    for (auto&& elt : sortSpec) {
-        if (elt.isNumber()) {
-            btreeBob.append(elt);
-            _patternPartTypes.push_back(SortPatternPartType::kFieldPath);
-        } else {
-            // If this field of the sort pattern is non-numeric, we expect it to be a text-score
-            // meta sort.
-            invariant(elt.type() == BSONType::Object);
-            invariant(elt.embeddedObject().nFields() == 1);
-            auto metaElem = elt.embeddedObject().firstElement();
-            invariant(metaElem.fieldNameStringData() == "$meta"_sd);
-            if (metaElem.valueStringData() == "textScore"_sd) {
-                _patternPartTypes.push_back(SortPatternPartType::kMetaTextScore);
-            } else if (metaElem.valueStringData() == "randVal"_sd) {
-                _patternPartTypes.push_back(SortPatternPartType::kMetaRandVal);
-            } else if (metaElem.valueStringData() == "searchScore"_sd) {
-                uasserted(31218, "$meta sort by 'searchScore' metadata is not supported");
-            } else if (metaElem.valueStringData() == "searchHighlights"_sd) {
-                uasserted(31219, "$meta sort by 'searchHighlights' metadata is not supported");
-            } else {
-                uasserted(31138, "Illegal $meta sort: " + metaElem.valueStringData());
-            }
-            _sortHasMeta = true;
+    for (auto&& part : _sortPattern) {
+        if (part.fieldPath) {
+            btreeBob.append(part.fieldPath->fullPath(), part.isAscending ? 1 : -1);
+            ++nFields;
         }
     }
 
     // The fake index key pattern used to generate Btree keys.
     _sortSpecWithoutMeta = btreeBob.obj();
+    _sortHasMeta = nFields < _sortPattern.size();
 
-    // If we're just sorting by meta, don't bother with all the key stuff.
+    // If we're just sorting by meta, don't bother creating an index key generator.
     if (_sortSpecWithoutMeta.isEmpty()) {
         return;
     }
@@ -77,11 +60,11 @@ SortKeyGenerator::SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterf
     // We'll need to treat arrays as if we were to create an index over them. that is, we may need
     // to unnest the first level and consider each array element to decide the sort order. In order
     // to do this, we make a BtreeKeyGenerator.
+    std::vector<BSONElement> fixed(nFields);
     std::vector<const char*> fieldNames;
-    std::vector<BSONElement> fixed;
-    for (auto&& patternElt : _sortSpecWithoutMeta) {
-        fieldNames.push_back(patternElt.fieldName());
-        fixed.push_back(BSONElement());
+    fieldNames.reserve(fixed.size());
+    for (auto&& elem : _sortSpecWithoutMeta) {
+        fieldNames.push_back(elem.fieldName());
     }
 
     constexpr bool isSparse = false;
@@ -105,10 +88,10 @@ StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromIndexKey(const WorkingSetMem
     invariant(!_sortHasMeta);
 
     BSONObjBuilder objBuilder;
-    for (BSONElement specElt : _sortSpecWithoutMeta) {
-        invariant(specElt.isNumber());
+    for (auto&& elem : _sortSpecWithoutMeta) {
         BSONElement sortKeyElt;
-        invariant(member.getFieldDotted(specElt.fieldName(), &sortKeyElt));
+        invariant(elem.isNumber());
+        invariant(member.getFieldDotted(elem.fieldName(), &sortKeyElt));
         // If we were to call 'collationAwareIndexKeyAppend' with a non-simple collation and a
         // 'sortKeyElt' representing a collated index key we would incorrectly encode for the
         // collation twice. This is not currently possible as the query planner will ensure that
@@ -139,18 +122,19 @@ StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromDocument(const BSONObj& obj,
 
     // Merge metadata into the key.
     BSONObjIterator sortKeyIt(sortKeyNoMetadata.getValue());
-    for (auto type : _patternPartTypes) {
-        switch (type) {
-            case SortPatternPartType::kFieldPath: {
-                invariant(sortKeyIt.more());
-                mergedKeyBob.append(sortKeyIt.next());
-                continue;
-            }
-            case SortPatternPartType::kMetaTextScore: {
+    for (auto& part : _sortPattern) {
+        if (part.fieldPath) {
+            invariant(sortKeyIt.more());
+            mergedKeyBob.append(sortKeyIt.next());
+            continue;
+        }
+        invariant(part.expression);
+        switch (part.expression->getMetaType()) {
+            case ExpressionMeta::MetaType::TEXT_SCORE: {
                 mergedKeyBob.append("", metadata->textScore);
                 continue;
             }
-            case SortPatternPartType::kMetaRandVal: {
+            case ExpressionMeta::MetaType::RAND_VAL: {
                 mergedKeyBob.append("", metadata->randVal);
                 continue;
             }
diff --git a/src/mongo/db/index/sort_key_generator.h b/src/mongo/db/index/sort_key_generator.h
index 64afec01cb8..2ddc186f340 100644
--- a/src/mongo/db/index/sort_key_generator.h
+++ b/src/mongo/db/index/sort_key_generator.h
@@ -34,6 +34,7 @@
 #include "mongo/db/index/btree_key_generator.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/query/collation/collator_interface.h"
+#include "mongo/db/query/sort_pattern.h"
 
 namespace mongo {
 
@@ -50,11 +51,11 @@ public:
     };
 
     /**
-     * Constructs a sort key generator which will generate keys for sort pattern 'sortSpec'. The
+     * Constructs a sort key generator which will generate keys for sort pattern 'sortPattern'. The
      * keys will incorporate the collation given by 'collator', and thus when actually compared to
      * one another should use the simple collation.
      */
-    SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterface* collator);
+    SortKeyGenerator(SortPattern sortPattern, const CollatorInterface* collator);
 
     /**
      * Returns the key which should be used to sort the WorkingSetMember, or a non-OK status if no
@@ -77,14 +78,6 @@ public:
     StatusWith<BSONObj> getSortKeyFromDocument(const BSONObj& obj, const Metadata*) const;
 
 private:
-    // Describes whether a component of the sort pattern is a field path (e.g. sort by "a.b"), or
-    // else describes the type of $meta sort.
-    enum class SortPatternPartType {
-        kFieldPath,
-        kMetaTextScore,
-        kMetaRandVal,
-    };
-
     // Extracts the sort key from a WorkingSetMember which represents an index key. It is illegal to
     // call this if the working set member is not in RID_AND_IDX state. It is also illegal to call
     // this if the sort pattern has any $meta components.
@@ -97,14 +90,12 @@ private:
 
     const CollatorInterface* _collator = nullptr;
 
+    SortPattern _sortPattern;
+
     // The sort pattern with any $meta sort components stripped out, since the underlying index key
     // generator does not understand $meta sort.
     BSONObj _sortSpecWithoutMeta;
 
-    // For each element of the raw sort spec, describes whether the element is sorting by a field
-    // path or by a particular meta-sort.
-    std::vector<SortPatternPartType> _patternPartTypes;
-
     // If we're not sorting with a $meta value we can short-cut some work.
     bool _sortHasMeta = false;
 
diff --git a/src/mongo/db/index/sort_key_generator_test.cpp b/src/mongo/db/index/sort_key_generator_test.cpp
index 36d910349c0..701c8975fd4 100644
--- a/src/mongo/db/index/sort_key_generator_test.cpp
+++ b/src/mongo/db/index/sort_key_generator_test.cpp
@@ -33,6 +33,7 @@
 
 #include "mongo/bson/json.h"
 #include "mongo/db/index/sort_key_generator.h"
+#include "mongo/db/pipeline/expression_context_for_test.h"
 #include "mongo/db/query/collation/collator_interface_mock.h"
 #include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
@@ -40,15 +41,24 @@
 namespace mongo {
 namespace {
 
+// A method to create mock ExpressionContexts with a specified collation
+std::unique_ptr<SortKeyGenerator> makeSortKeyGen(const BSONObj& sortSpec,
+                                                 const CollatorInterface* collator) {
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContextForTest());
+    pExpCtx->setCollator(collator);
+    SortPattern sortPattern{sortSpec, pExpCtx};
+    return std::make_unique<SortKeyGenerator>(std::move(sortPattern), collator);
+}
+
 TEST(SortKeyGeneratorTest, ExtractNumberKeyForNonCompoundSortNonNested) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, a: 5}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 5));
 }
 
 TEST(SortKeyGeneratorTest, ExtractNumberKeyFromDocWithSeveralFields) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     auto sortKey =
         sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -56,7 +66,7 @@ TEST(SortKeyGeneratorTest, ExtractNumberKeyFromDocWithSeveralFields) {
 }
 
 TEST(SortKeyGeneratorTest, ExtractStringKeyNonCompoundNonNested) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -66,7 +76,7 @@ TEST(SortKeyGeneratorTest, ExtractStringKeyNonCompoundNonNested) {
 }
 
 TEST(SortKeyGeneratorTest, CompoundSortPattern) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1 << "b" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1 << "b" << 1), nullptr);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -74,7 +84,7 @@ TEST(SortKeyGeneratorTest, CompoundSortPattern) {
 }
 
 TEST(SortKeyGeneratorTest, CompoundSortPatternWithDottedPath) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("c.a" << 1 << "b" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("c.a" << 1 << "b" << 1), nullptr);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -82,7 +92,7 @@ TEST(SortKeyGeneratorTest, CompoundSortPatternWithDottedPath) {
 }
 
 TEST(SortKeyGeneratorTest, CompoundPatternLeadingFieldIsArray) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("c" << 1 << "b" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("c" << 1 << "b" << 1), nullptr);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, z: 'thing1', a: 99, c: [2, 4, 1], b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -91,7 +101,7 @@ TEST(SortKeyGeneratorTest, CompoundPatternLeadingFieldIsArray) {
 
 TEST(SortKeyGeneratorTest, ExtractStringSortKeyWithCollatorUsesComparisonKey) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -102,7 +112,7 @@ TEST(SortKeyGeneratorTest, ExtractStringSortKeyWithCollatorUsesComparisonKey) {
 
 TEST(SortKeyGeneratorTest, CollatorHasNoEffectWhenExtractingNonStringSortKey) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator);
     auto sortKey =
         sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -110,7 +120,7 @@ TEST(SortKeyGeneratorTest, CollatorHasNoEffectWhenExtractingNonStringSortKey) {
 }
 
 TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysChoosesCorrectKey) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << -1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << -1), nullptr);
     auto sortKey =
         sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, a: [1, 2, 3, 4]}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -119,7 +129,7 @@ TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysChoosesCorrectKey) {
 
 TEST(SortKeyGeneratorTest, EnsureSortKeyGenerationForArraysRespectsCollation) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, a: ['aaz', 'zza', 'yya', 'zzb']}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -129,7 +139,7 @@ TEST(SortKeyGeneratorTest, EnsureSortKeyGenerationForArraysRespectsCollation) {
 }
 
 TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysRespectsCompoundOrdering) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a.b" << 1 << "a.c" << -1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a.b" << 1 << "a.c" << -1), nullptr);
     auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, a: [{b: 1, c: 0}, {b: 0, c: 3}, {b: 0, c: 1}]}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
@@ -138,62 +148,59 @@ TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysRespectsCompoundOrdering) {
 
 DEATH_TEST(SortKeyGeneratorTest,
            SortPatternComponentWithStringIsFatal,
-           "Invariant failure elt.type() == BSONType::Object") {
-    MONGO_COMPILER_VARIABLE_UNUSED auto ignored = std::make_unique<SortKeyGenerator>(BSON("a"
-                                                                                          << "foo"),
-                                                                                     nullptr);
+           "Illegal key in $sort specification: a: \"foo\"") {
+    MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a"
+                                                                      << "foo"),
+                                                                 nullptr);
 }
 
 DEATH_TEST(SortKeyGeneratorTest,
            SortPatternComponentWhoseObjectHasMultipleKeysIsFatal,
-           "Invariant failure elt.embeddedObject().nFields() == 1") {
-    MONGO_COMPILER_VARIABLE_UNUSED auto ignored =
-        std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
-                                                            << "textScore"
-                                                            << "extra" << 1)),
-                                           nullptr);
+           "Cannot have additional keys in a $meta sort specification") {
+    MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a" << BSON("$meta"
+                                                                                  << "textScore"
+                                                                                  << "extra" << 1)),
+                                                                 nullptr);
 }
 
 DEATH_TEST(SortKeyGeneratorTest,
            SortPatternComponentWithNonMetaObjectSortIsFatal,
-           "Invariant failure metaElem.fieldNameStringData() == \"$meta\"_sd") {
-    MONGO_COMPILER_VARIABLE_UNUSED auto ignored =
-        std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$unknown"
-                                                            << "textScore")),
-                                           nullptr);
+           "$meta is the only expression supported by $sort right now") {
+    MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a" << BSON("$unknown"
+                                                                                  << "textScore")),
+                                                                 nullptr);
 }
 
 DEATH_TEST(SortKeyGeneratorTest,
            SortPatternComponentWithUnknownMetaKeywordIsFatal,
            "Illegal $meta sort: unknown") {
-    MONGO_COMPILER_VARIABLE_UNUSED auto ignored =
-        std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
-                                                            << "unknown")),
-                                           nullptr);
+    MONGO_COMPILER_VARIABLE_UNUSED auto ignored = makeSortKeyGen(BSON("a" << BSON("$meta"
+                                                                                  << "unknown")),
+                                                                 nullptr);
 }
 
 DEATH_TEST(SortKeyGeneratorTest,
            NoMetadataWhenPatternHasMetaTextScoreIsFatal,
            "Invariant failure metadata") {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
-                                                                          << "textScore")),
-                                                         nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta"
+                                                      << "textScore")),
+                                     nullptr);
     uassertStatusOK(sortKeyGen->getSortKeyFromDocument(BSONObj{}, nullptr).getStatus());
 }
 
 DEATH_TEST(SortKeyGeneratorTest,
            NoMetadataWhenPatternHasMetaRandValIsFatal,
            "Invariant failure metadata") {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
-                                                                          << "randVal")),
-                                                         nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta"
+                                                      << "randVal")),
+                                     nullptr);
     uassertStatusOK(sortKeyGen->getSortKeyFromDocument(BSONObj{}, nullptr).getStatus());
 }
 
 TEST(SortKeyGeneratorTest, CanGenerateKeysForTextScoreMetaSort) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
-                                                                          << "textScore")),
-                                                         nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta"
+                                                      << "textScore")),
+                                     nullptr);
     SortKeyGenerator::Metadata metadata;
     metadata.textScore = 1.5;
     auto sortKey = sortKeyGen->getSortKeyFromDocument(BSONObj{}, &metadata);
@@ -202,9 +209,9 @@ TEST(SortKeyGeneratorTest, CanGenerateKeysForTextScoreMetaSort) {
 }
 
 TEST(SortKeyGeneratorTest, CanGenerateKeysForRandValMetaSort) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
-                                                                          << "randVal")),
-                                                         nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << BSON("$meta"
+                                                      << "randVal")),
+                                     nullptr);
     SortKeyGenerator::Metadata metadata;
     metadata.randVal = 0.3;
     auto sortKey = sortKeyGen->getSortKeyFromDocument(BSONObj{}, &metadata);
@@ -215,7 +222,7 @@ TEST(SortKeyGeneratorTest, CanGenerateKeysForRandValMetaSort) {
 TEST(SortKeyGeneratorTest, CanGenerateKeysForCompoundMetaSort) {
     BSONObj pattern = fromjson(
         "{a: 1, b: {$meta: 'randVal'}, c: {$meta: 'textScore'}, d: -1, e: {$meta: 'textScore'}}");
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr);
+    auto sortKeyGen = makeSortKeyGen(pattern, nullptr);
     SortKeyGenerator::Metadata metadata;
     metadata.randVal = 0.3;
     metadata.textScore = 1.5;
@@ -258,7 +265,7 @@ private:
 };
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithObj) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     setRecordIdAndObj(BSON("x" << 1 << "a" << 2 << "y" << 3));
     auto sortKey = sortKeyGen->getSortKey(member());
     ASSERT_OK(sortKey);
@@ -266,7 +273,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithObj)
 }
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithOwnedObj) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3));
     auto sortKey = sortKeyGen->getSortKey(member());
     ASSERT_OK(sortKey);
@@ -275,7 +282,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithOwne
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateKeyFromWSMForTextScoreMetaSort) {
     BSONObj pattern = fromjson("{a: 1, b: {$meta: 'textScore'}, c: -1}}");
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr);
+    auto sortKeyGen = makeSortKeyGen(pattern, nullptr);
     setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3 << "c" << BSON_ARRAY(4 << 5 << 6)));
     member().metadata().setTextScore(9.9);
     auto sortKey = sortKeyGen->getSortKey(member());
@@ -284,7 +291,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateKeyFromWSMForTextScoreMetaSort
 }
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyState) {
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << 2 << "" << 3));
     auto sortKey = sortKeyGen->getSortKey(member());
     ASSERT_OK(sortKey);
@@ -293,7 +300,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyState)
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyStateWithCollator) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
+    auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), &collator);
     setRecordIdAndIdx(BSON("a" << 1 << "b" << 1),
                       BSON(""
                            << "string1"
@@ -310,7 +317,7 @@ DEATH_TEST_F(SortKeyGeneratorWorkingSetTest,
              DeathOnAttemptToGetSortKeyFromIndexKeyWithMetadata,
              "Invariant failure !_sortHasMeta") {
     BSONObj pattern = fromjson("{z: {$meta: 'textScore'}}");
-    auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr);
+    auto sortKeyGen = makeSortKeyGen(pattern, nullptr);
     setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << 2 << "" << 3));
     member().metadata().setTextScore(9.9);
     MONGO_COMPILER_VARIABLE_UNUSED auto ignored = sortKeyGen->getSortKey(member());
diff --git a/src/mongo/db/pipeline/document_source_sort.cpp b/src/mongo/db/pipeline/document_source_sort.cpp
index fe32168c170..09afc8fc2e6 100644
--- a/src/mongo/db/pipeline/document_source_sort.cpp
+++ b/src/mongo/db/pipeline/document_source_sort.cpp
@@ -112,7 +112,7 @@ DocumentSourceSort::DocumentSourceSort(const boost::intrusive_ptr<ExpressionCont
                      pExpCtx->allowDiskUse}),
       // The SortKeyGenerator expects the expressions to be serialized in order to detect a sort
       // by a metadata field.
-      _sortKeyGen({sortOrder, pExpCtx->getCollator()}) {
+      _sortKeyGen({{sortOrder, pExpCtx}, pExpCtx->getCollator()}) {
     uassert(15976,
             "$sort stage must have at least one sort key",
             !_sortExecutor->sortPattern().empty());
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index ab70a4d4319..641d35d5aa4 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1642,6 +1642,13 @@ private:
 
 class ExpressionMeta final : public Expression {
 public:
+    enum MetaType {
+        TEXT_SCORE,
+        RAND_VAL,
+        SEARCH_SCORE,
+        SEARCH_HIGHLIGHTS,
+    };
+
     Value serialize(bool explain) const final;
     Value evaluate(const Document& root, Variables* variables) const final;
 
@@ -1654,17 +1661,14 @@ public:
         return visitor->visit(this);
     }
 
+    MetaType getMetaType() {
+        return _metaType;
+    }
+
 protected:
     void _doAddDependencies(DepsTracker* deps) const final;
 
 private:
-    enum MetaType {
-        TEXT_SCORE,
-        RAND_VAL,
-        SEARCH_SCORE,
-        SEARCH_HIGHLIGHTS,
-    };
-
     ExpressionMeta(const boost::intrusive_ptr<ExpressionContext>& expCtx, MetaType metaType);
 
     MetaType _metaType;
diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index 03424a838d6..e0a164bd2a2 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -205,7 +205,7 @@ StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::canonicalize(
     // Make the CQ we'll hopefully return.
     std::unique_ptr<CanonicalQuery> cq(new CanonicalQuery());
     Status initStatus = cq->init(opCtx,
-                                 nullptr,  // no expression context
+                                 baseQuery.getExpCtx(),
                                  std::move(qr),
                                  baseQuery.canHaveNoopMatchNodes(),
                                  root->shallowClone(),
@@ -261,6 +261,9 @@ void CanonicalQuery::setCollator(std::unique_ptr<CollatorInterface> collator) {
     // the object owned by '_collator'. We must associate the match expression tree with the new
     // value of '_collator'.
     _root->setCollator(_collator.get());
+
+    // In a similar vein, we must give the ExpressionContext the same collator.
+    _expCtx->setCollator(_collator.get());
 }
 
 // static
diff --git a/src/mongo/db/query/canonical_query.h b/src/mongo/db/query/canonical_query.h
index 366265add3c..bbd7103707d 100644
--- a/src/mongo/db/query/canonical_query.h
+++ b/src/mongo/db/query/canonical_query.h
@@ -181,6 +181,10 @@ public:
         return _canHaveNoopMatchNodes;
     }
 
+    const boost::intrusive_ptr<ExpressionContext>& getExpCtx() const {
+        return _expCtx;
+    }
+
 private:
     // You must go through canonicalize to create a CanonicalQuery.
     CanonicalQuery() {}
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 178821fb880..7e8c0d1e9b4 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -395,11 +395,10 @@ StatusWith<PrepareExecutionResult> prepareExecution(OperationContext* opCtx,
             // Add a SortKeyGeneratorStage if there is a $meta sortKey projection.
             if (canonicalQuery->getProj()->wantSortKey()) {
                 root = std::make_unique<SortKeyGeneratorStage>(
-                    opCtx,
+                    canonicalQuery->getExpCtx(),
                     root.release(),
                     ws,
-                    canonicalQuery->getQueryRequest().getSort(),
-                    canonicalQuery->getCollator());
+                    canonicalQuery->getQueryRequest().getSort());
             }
 
             // Stuff the right data into the params depending on what proj impl we use.
diff --git a/src/mongo/db/query/sort_pattern.cpp b/src/mongo/db/query/sort_pattern.cpp
index 17df5adcca7..1964164cf5e 100644
--- a/src/mongo/db/query/sort_pattern.cpp
+++ b/src/mongo/db/query/sort_pattern.cpp
@@ -51,7 +51,22 @@ SortPattern::SortPattern(const BSONObj& obj,
                     metaDoc.nFields() == 1);
 
             VariablesParseState vps = pExpCtx->variablesParseState;
-            patternPart.expression = ExpressionMeta::parse(pExpCtx, metaDoc.firstElement(), vps);
+            BSONElement metaElem = metaDoc.firstElement();
+
+            if (metaElem.valueStringData() == "textScore"_sd) {
+                // Valid meta sort. Just fall through.
+            } else if (metaElem.valueStringData() == "randVal"_sd) {
+                // Valid meta sort. Just fall through.
+            } else if (metaElem.valueStringData() == "searchScore"_sd) {
+                uasserted(31218, "$meta sort by 'searchScore' metadata is not supported");
+            } else if (metaElem.valueStringData() == "searchHighlights"_sd) {
+                uasserted(31219, "$meta sort by 'searchHighlights' metadata is not supported");
+            } else {
+                uasserted(31138,
+                          str::stream() << "Illegal $meta sort: " << metaElem.valueStringData());
+            }
+            patternPart.expression =
+                static_cast<ExpressionMeta*>(ExpressionMeta::parse(pExpCtx, metaElem, vps).get());
 
             // If sorting by textScore, sort highest scores first. If sorting by randVal, order
             // doesn't matter, so just always use descending.
diff --git a/src/mongo/db/query/sort_pattern.h b/src/mongo/db/query/sort_pattern.h
index 5dbb00c7ca2..d818bd82f38 100644
--- a/src/mongo/db/query/sort_pattern.h
+++ b/src/mongo/db/query/sort_pattern.h
@@ -49,7 +49,7 @@ public:
     struct SortPatternPart {
         bool isAscending = true;
         boost::optional<FieldPath> fieldPath;
-        boost::intrusive_ptr<Expression> expression;
+        boost::intrusive_ptr<ExpressionMeta> expression;
     };
 
     SortPattern(const BSONObj&, const boost::intrusive_ptr<ExpressionContext>&);
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index dcba367378b..382e3652354 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -142,8 +142,7 @@ PlanStage* buildStages(OperationContext* opCtx,
             if (nullptr == childStage) {
                 return nullptr;
             }
-            return new SortKeyGeneratorStage(
-                opCtx, childStage, ws, keyGenNode->sortSpec, cq.getCollator());
+            return new SortKeyGeneratorStage(cq.getExpCtx(), childStage, ws, keyGenNode->sortSpec);
         }
         case STAGE_PROJECTION_DEFAULT: {
             auto pn = static_cast<const ProjectionNodeDefault*>(root);
diff --git a/src/mongo/dbtests/query_stage_ensure_sorted.cpp b/src/mongo/dbtests/query_stage_ensure_sorted.cpp
index 49e9383442f..20d5750c958 100644
--- a/src/mongo/dbtests/query_stage_ensure_sorted.cpp
+++ b/src/mongo/dbtests/query_stage_ensure_sorted.cpp
@@ -78,10 +78,15 @@ public:
             queuedDataStage->pushBack(id);
         }
 
+        // Create a mock ExpressionContext.
+        boost::intrusive_ptr<ExpressionContext> pExpCtx(
+            new ExpressionContext(opCtx.get(), collator));
+        pExpCtx->setCollator(collator);
+
         // Initialization.
         BSONObj pattern = fromjson(patternStr);
         auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-            opCtx.get(), queuedDataStage.release(), &ws, pattern, collator);
+            pExpCtx, queuedDataStage.release(), &ws, pattern);
         EnsureSortedStage ess(opCtx.get(), pattern, &ws, sortKeyGen.release());
         WorkingSetID id = WorkingSet::INVALID_ID;
         PlanStage::StageState state = PlanStage::NEED_TIME;
@@ -115,10 +120,13 @@ protected:
 TEST_F(QueryStageEnsureSortedTest, EnsureSortedEmptyWorkingSet) {
     auto opCtx = _serviceContext.makeOperationContext();
 
+    // Create a mock ExpressionContext.
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), nullptr));
+
     WorkingSet ws;
     auto queuedDataStage = std::make_unique<QueuedDataStage>(opCtx.get(), &ws);
-    auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-        opCtx.get(), queuedDataStage.release(), &ws, BSONObj(), nullptr);
+    auto sortKeyGen =
+        std::make_unique<SortKeyGeneratorStage>(pExpCtx, queuedDataStage.release(), &ws, BSONObj());
     EnsureSortedStage ess(opCtx.get(), BSONObj(), &ws, sortKeyGen.release());
 
     WorkingSetID id = WorkingSet::INVALID_ID;
diff --git a/src/mongo/dbtests/query_stage_sort.cpp b/src/mongo/dbtests/query_stage_sort.cpp
index 5b855933793..0919fdea1ad 100644
--- a/src/mongo/dbtests/query_stage_sort.cpp
+++ b/src/mongo/dbtests/query_stage_sort.cpp
@@ -120,7 +120,7 @@ public:
         params.limit = limit();
 
         auto keyGenStage = std::make_unique<SortKeyGeneratorStage>(
-            &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr);
+            _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern);
 
         auto ss = std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release());
 
@@ -158,7 +158,7 @@ public:
         params.limit = limit();
 
         auto keyGenStage = std::make_unique<SortKeyGeneratorStage>(
-            &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr);
+            _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern);
 
         auto sortStage =
             std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release());
@@ -227,6 +227,7 @@ public:
 protected:
     const ServiceContext::UniqueOperationContext _txnPtr = cc().makeOperationContext();
     OperationContext& _opCtx = *_txnPtr;
+    boost::intrusive_ptr<ExpressionContext> _pExpCtx = new ExpressionContext(&_opCtx, nullptr);
     DBDirectClient _client;
 };
 
@@ -558,7 +559,7 @@ public:
         params.pattern = BSON("b" << -1 << "c" << 1 << "a" << 1);
 
         auto keyGenStage = std::make_unique<SortKeyGeneratorStage>(
-            &_opCtx, queuedDataStage.release(), ws.get(), params.pattern, nullptr);
+            _pExpCtx, queuedDataStage.release(), ws.get(), params.pattern);
 
         auto sortStage =
             std::make_unique<SortStage>(&_opCtx, params, ws.get(), keyGenStage.release());
diff --git a/src/mongo/dbtests/query_stage_sort_key_generator.cpp b/src/mongo/dbtests/query_stage_sort_key_generator.cpp
index daa667313e4..8675a3e7a51 100644
--- a/src/mongo/dbtests/query_stage_sort_key_generator.cpp
+++ b/src/mongo/dbtests/query_stage_sort_key_generator.cpp
@@ -65,6 +65,7 @@ BSONObj extractKeyFromKeyGenStage(SortKeyGeneratorStage* sortKeyGen, WorkingSet*
 BSONObj extractSortKey(const char* sortSpec, const char* doc, const CollatorInterface* collator) {
     QueryTestServiceContext serviceContext;
     auto opCtx = serviceContext.makeOperationContext();
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), collator));
 
     WorkingSet workingSet;
 
@@ -77,7 +78,7 @@ BSONObj extractSortKey(const char* sortSpec, const char* doc, const CollatorInte
 
     BSONObj sortPattern = fromjson(sortSpec);
     SortKeyGeneratorStage sortKeyGen{
-        opCtx.get(), mockStage.release(), &workingSet, std::move(sortPattern), collator};
+        pExpCtx, mockStage.release(), &workingSet, std::move(sortPattern)};
     return extractKeyFromKeyGenStage(&sortKeyGen, &workingSet);
 }
 
@@ -93,6 +94,7 @@ BSONObj extractSortKeyCovered(const char* sortSpec,
                               const CollatorInterface* collator) {
     QueryTestServiceContext serviceContext;
     auto opCtx = serviceContext.makeOperationContext();
+    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(opCtx.get(), collator));
 
     WorkingSet workingSet;
 
@@ -105,7 +107,7 @@ BSONObj extractSortKeyCovered(const char* sortSpec,
 
     BSONObj sortPattern = fromjson(sortSpec);
     SortKeyGeneratorStage sortKeyGen{
-        opCtx.get(), mockStage.release(), &workingSet, std::move(sortPattern), collator};
+        pExpCtx, mockStage.release(), &workingSet, std::move(sortPattern)};
     return extractKeyFromKeyGenStage(&sortKeyGen, &workingSet);
 }