SERVER-42161 Move sort key generation code for RID_AND_IDX into SortKeyGenerator.

This is a pure refactor which helps consolidate all sort key
generation logic into SortKeyGenerator. This is a step along
the way towards unifying SortStage and DocumentSourceSort.

8 files changed, 188 insertions, 93 deletions

diff --git a/src/mongo/db/exec/sort_key_generator.cpp b/src/mongo/db/exec/sort_key_generator.cpp
index 5fb3fd905ee..1aa3abf8660 100644
--- a/src/mongo/db/exec/sort_key_generator.cpp
+++ b/src/mongo/db/exec/sort_key_generator.cpp
@@ -43,7 +43,6 @@
 #include "mongo/db/exec/working_set_common.h"
 #include "mongo/db/exec/working_set_computed_data.h"
 #include "mongo/db/matcher/extensions_callback_noop.h"
-#include "mongo/db/query/collation/collation_index_key.h"
 #include "mongo/db/query/collation/collator_interface.h"
 #include "mongo/util/log.h"
 
@@ -56,7 +55,7 @@ SortKeyGeneratorStage::SortKeyGeneratorStage(OperationContext* opCtx,
                                              WorkingSet* ws,
                                              const BSONObj& sortSpecObj,
                                              const CollatorInterface* collator)
-    : PlanStage(kStageType, opCtx), _ws(ws), _sortSpec(sortSpecObj), _collator(collator) {
+    : PlanStage(kStageType, opCtx), _ws(ws), _sortKeyGen(sortSpecObj, collator) {
     _children.emplace_back(child);
 }
 
@@ -65,28 +64,11 @@ bool SortKeyGeneratorStage::isEOF() {
 }
 
 PlanStage::StageState SortKeyGeneratorStage::doWork(WorkingSetID* out) {
-    if (!_sortKeyGen) {
-        _sortKeyGen = std::make_unique<SortKeyGenerator>(_sortSpec, _collator);
-        return PlanStage::NEED_TIME;
-    }
-
     auto stageState = child()->work(out);
     if (stageState == PlanStage::ADVANCED) {
         WorkingSetMember* member = _ws->get(*out);
 
-        StatusWith<BSONObj> sortKey = BSONObj();
-        if (member->hasObj()) {
-            SortKeyGenerator::Metadata metadata;
-            if (_sortKeyGen->sortHasMeta() && member->hasComputed(WSM_COMPUTED_TEXT_SCORE)) {
-                auto scoreData = static_cast<const TextScoreComputedData*>(
-                    member->getComputed(WSM_COMPUTED_TEXT_SCORE));
-                metadata.textScore = scoreData->getScore();
-            }
-            sortKey = _sortKeyGen->getSortKey(member->obj.value(), &metadata);
-        } else {
-            sortKey = getSortKeyFromIndexKey(*member);
-        }
-
+        auto sortKey = _sortKeyGen.getSortKey(*member);
         if (!sortKey.isOK()) {
             *out = WorkingSetCommon::allocateStatusMember(_ws, sortKey.getStatus());
             return PlanStage::FAILURE;
@@ -116,24 +98,4 @@ const SpecificStats* SortKeyGeneratorStage::getSpecificStats() const {
     return nullptr;
 }
 
-StatusWith<BSONObj> SortKeyGeneratorStage::getSortKeyFromIndexKey(
-    const WorkingSetMember& member) const {
-    invariant(member.getState() == WorkingSetMember::RID_AND_IDX);
-    invariant(!_sortKeyGen->sortHasMeta());
-
-    BSONObjBuilder objBuilder;
-    for (BSONElement specElt : _sortSpec) {
-        invariant(specElt.isNumber());
-        BSONElement sortKeyElt;
-        invariant(member.getFieldDotted(specElt.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
-        // the plan fetches the data before sort key generation in the case where the index has a
-        // non-simple collation.
-        CollationIndexKey::collationAwareIndexKeyAppend(sortKeyElt, _collator, &objBuilder);
-    }
-    return objBuilder.obj();
-}
-
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sort_key_generator.h b/src/mongo/db/exec/sort_key_generator.h
index b85a54bc18a..b5208ee78fb 100644
--- a/src/mongo/db/exec/sort_key_generator.h
+++ b/src/mongo/db/exec/sort_key_generator.h
@@ -71,16 +71,9 @@ protected:
     StageState doWork(WorkingSetID* out) final;
 
 private:
-    StatusWith<BSONObj> getSortKeyFromIndexKey(const WorkingSetMember& member) const;
-
     WorkingSet* const _ws;
 
-    // The raw sort pattern as expressed by the user.
-    const BSONObj _sortSpec;
-
-    const CollatorInterface* _collator;
-
-    std::unique_ptr<SortKeyGenerator> _sortKeyGen;
+    SortKeyGenerator _sortKeyGen;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sort_test.cpp b/src/mongo/db/exec/sort_test.cpp
index e693bf6b412..66d6f3c1d86 100644
--- a/src/mongo/db/exec/sort_test.cpp
+++ b/src/mongo/db/exec/sort_test.cpp
@@ -168,13 +168,9 @@ TEST_F(SortStageTest, SortEmptyWorkingSet) {
     // Check initial EOF state.
     ASSERT_FALSE(sort.isEOF());
 
-    // First call to work() initializes sort key generator.
+    // First call to work() sorts data in vector.
     WorkingSetID id = WorkingSet::INVALID_ID;
-    PlanStage::StageState state = sort.work(&id);
-    ASSERT_EQUALS(state, PlanStage::NEED_TIME);
-
-    // Second call to work() sorts data in vector.
-    state = sort.work(&id);
+    auto state = sort.work(&id);
     ASSERT_EQUALS(state, PlanStage::NEED_TIME);
 
     // Finally we hit EOF.
diff --git a/src/mongo/db/index/SConscript b/src/mongo/db/index/SConscript
index 6247ba89530..07784425da7 100644
--- a/src/mongo/db/index/SConscript
+++ b/src/mongo/db/index/SConscript
@@ -46,6 +46,7 @@ env.Library(
         LIBDEPS=[
             '$BUILD_DIR/mongo/base',
             '$BUILD_DIR/mongo/db/bson/dotted_path_support',
+            '$BUILD_DIR/mongo/db/exec/working_set',
             '$BUILD_DIR/mongo/db/fts/base_fts',
             '$BUILD_DIR/mongo/db/geo/geoparser',
             '$BUILD_DIR/mongo/db/index_names',
diff --git a/src/mongo/db/index/sort_key_generator.cpp b/src/mongo/db/index/sort_key_generator.cpp
index 768e2258e2e..a9fd65038fb 100644
--- a/src/mongo/db/index/sort_key_generator.cpp
+++ b/src/mongo/db/index/sort_key_generator.cpp
@@ -32,6 +32,8 @@
 #include "mongo/db/index/sort_key_generator.h"
 
 #include "mongo/bson/bsonobj_comparator.h"
+#include "mongo/db/exec/working_set_computed_data.h"
+#include "mongo/db/query/collation/collation_index_key.h"
 
 namespace mongo {
 
@@ -82,26 +84,59 @@ SortKeyGenerator::SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterf
     _indexKeyGen = std::make_unique<BtreeKeyGenerator>(fieldNames, fixed, isSparse, _collator);
 }
 
-StatusWith<BSONObj> SortKeyGenerator::getSortKey(const BSONObj& obj,
-                                                 const Metadata* metadata) const {
+StatusWith<BSONObj> SortKeyGenerator::getSortKey(const WorkingSetMember& wsm) const {
+    if (wsm.hasObj()) {
+        SortKeyGenerator::Metadata metadata;
+        if (_sortHasMeta && wsm.hasComputed(WSM_COMPUTED_TEXT_SCORE)) {
+            auto scoreData =
+                static_cast<const TextScoreComputedData*>(wsm.getComputed(WSM_COMPUTED_TEXT_SCORE));
+            metadata.textScore = scoreData->getScore();
+        }
+        return getSortKeyFromDocument(wsm.obj.value(), &metadata);
+    }
+
+    return getSortKeyFromIndexKey(wsm);
+}
+
+StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromIndexKey(const WorkingSetMember& member) const {
+    invariant(member.getState() == WorkingSetMember::RID_AND_IDX);
+    invariant(!_sortHasMeta);
+
+    BSONObjBuilder objBuilder;
+    for (BSONElement specElt : _sortSpecWithoutMeta) {
+        invariant(specElt.isNumber());
+        BSONElement sortKeyElt;
+        invariant(member.getFieldDotted(specElt.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
+        // the plan fetches the data before sort key generation in the case where the index has a
+        // non-simple collation.
+        CollationIndexKey::collationAwareIndexKeyAppend(sortKeyElt, _collator, &objBuilder);
+    }
+    return objBuilder.obj();
+}
+
+StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromDocument(const BSONObj& obj,
+                                                             const Metadata* metadata) const {
     if (_sortHasMeta) {
         invariant(metadata);
     }
 
-    auto indexKey = getIndexKey(obj);
-    if (!indexKey.isOK()) {
-        return indexKey;
+    auto sortKeyNoMetadata = getSortKeyFromDocumentWithoutMetadata(obj);
+    if (!sortKeyNoMetadata.isOK()) {
+        return sortKeyNoMetadata;
     }
 
     if (!_sortHasMeta) {
         // We don't have to worry about $meta sort, so the index key becomes the sort key.
-        return indexKey;
+        return sortKeyNoMetadata;
     }
 
     BSONObjBuilder mergedKeyBob;
 
     // Merge metadata into the key.
-    BSONObjIterator sortKeyIt(indexKey.getValue());
+    BSONObjIterator sortKeyIt(sortKeyNoMetadata.getValue());
     for (auto type : _patternPartTypes) {
         switch (type) {
             case SortPatternPartType::kFieldPath: {
@@ -127,7 +162,8 @@ StatusWith<BSONObj> SortKeyGenerator::getSortKey(const BSONObj& obj,
     return mergedKeyBob.obj();
 }
 
-StatusWith<BSONObj> SortKeyGenerator::getIndexKey(const BSONObj& obj) const {
+StatusWith<BSONObj> SortKeyGenerator::getSortKeyFromDocumentWithoutMetadata(
+    const BSONObj& obj) const {
     // Not sorting by anything in the key, just bail out early.
     if (_sortSpecWithoutMeta.isEmpty()) {
         return BSONObj();
diff --git a/src/mongo/db/index/sort_key_generator.h b/src/mongo/db/index/sort_key_generator.h
index 4775358f1e1..64afec01cb8 100644
--- a/src/mongo/db/index/sort_key_generator.h
+++ b/src/mongo/db/index/sort_key_generator.h
@@ -30,6 +30,7 @@
 #pragma once
 
 #include "mongo/bson/bsonobj.h"
+#include "mongo/db/exec/working_set.h"
 #include "mongo/db/index/btree_key_generator.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/query/collation/collator_interface.h"
@@ -56,6 +57,16 @@ public:
     SortKeyGenerator(const BSONObj& sortSpec, const CollatorInterface* collator);
 
     /**
+     * Returns the key which should be used to sort the WorkingSetMember, or a non-OK status if no
+     * key could be generated. The WorkingSetMember may represent either an index key, or a document
+     * (owned or unowned) that has been fetched from the collection.
+     *
+     * If the sort pattern contains a $meta sort (e.g. sort by "textScore" or "randVal"), then the
+     * necessary metadata is obtained from the WorkingSetMember.
+     */
+    StatusWith<BSONObj> getSortKey(const WorkingSetMember&) const;
+
+    /**
      * Returns the key which should be used to sort 'obj', or a non-OK status if no key could be
      * generated.
      *
@@ -63,14 +74,7 @@ public:
      * a $meta sort (i.e. if sortHasMeta() is true). These values are filled in at the corresponding
      * positions in the sort key.
      */
-    StatusWith<BSONObj> getSortKey(const BSONObj& obj, const Metadata*) const;
-
-    /**
-     * Returns true if the sort pattern for this sort key generator includes a $meta sort.
-     */
-    bool sortHasMeta() const {
-        return _sortHasMeta;
-    }
+    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
@@ -81,7 +85,15 @@ private:
         kMetaRandVal,
     };
 
-    StatusWith<BSONObj> getIndexKey(const BSONObj& obj) const;
+    // 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.
+    StatusWith<BSONObj> getSortKeyFromIndexKey(const WorkingSetMember& member) const;
+
+    // Extracts the sort key from 'obj', using '_sortSpecWithoutMeta' and thus ignoring any $meta
+    // sort components of the sort pattern. The caller is responsible for augmenting this key with
+    // the appropriate metadata if '_sortHasMeta' is true.
+    StatusWith<BSONObj> getSortKeyFromDocumentWithoutMetadata(const BSONObj& obj) const;
 
     const CollatorInterface* _collator = nullptr;
 
diff --git a/src/mongo/db/index/sort_key_generator_test.cpp b/src/mongo/db/index/sort_key_generator_test.cpp
index b129d99ad64..a533a6cfbff 100644
--- a/src/mongo/db/index/sort_key_generator_test.cpp
+++ b/src/mongo/db/index/sort_key_generator_test.cpp
@@ -32,6 +32,7 @@
 #include <memory>
 
 #include "mongo/bson/json.h"
+#include "mongo/db/exec/working_set_computed_data.h"
 #include "mongo/db/index/sort_key_generator.h"
 #include "mongo/db/query/collation/collator_interface_mock.h"
 #include "mongo/unittest/death_test.h"
@@ -42,22 +43,23 @@ namespace {
 
 TEST(SortKeyGeneratorTest, ExtractNumberKeyForNonCompoundSortNonNested) {
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
-    auto sortKey = sortKeyGen->getSortKey(fromjson("{_id: 0, a: 5}"), 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 sortKey = sortKeyGen->getSortKey(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr);
+    auto sortKey =
+        sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 6));
 }
 
 TEST(SortKeyGeneratorTest, ExtractStringKeyNonCompoundNonNested) {
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
-    auto sortKey =
-        sortKeyGen->getSortKey(fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(
+        fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(),
                       BSON(""
@@ -66,23 +68,23 @@ TEST(SortKeyGeneratorTest, ExtractStringKeyNonCompoundNonNested) {
 
 TEST(SortKeyGeneratorTest, CompoundSortPattern) {
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1 << "b" << 1), nullptr);
-    auto sortKey =
-        sortKeyGen->getSortKey(fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(
+        fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 99 << "" << 16));
 }
 
 TEST(SortKeyGeneratorTest, CompoundSortPatternWithDottedPath) {
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("c.a" << 1 << "b" << 1), nullptr);
-    auto sortKey =
-        sortKeyGen->getSortKey(fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(
+        fromjson("{_id: 0, z: 'thing1', a: 99, c: {a: 4}, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 4 << "" << 16));
 }
 
 TEST(SortKeyGeneratorTest, CompoundPatternLeadingFieldIsArray) {
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("c" << 1 << "b" << 1), nullptr);
-    auto sortKey = sortKeyGen->getSortKey(
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(
         fromjson("{_id: 0, z: 'thing1', a: 99, c: [2, 4, 1], b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 1 << "" << 16));
@@ -91,8 +93,8 @@ TEST(SortKeyGeneratorTest, CompoundPatternLeadingFieldIsArray) {
 TEST(SortKeyGeneratorTest, ExtractStringSortKeyWithCollatorUsesComparisonKey) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
-    auto sortKey =
-        sortKeyGen->getSortKey(fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(
+        fromjson("{_id: 0, z: 'thing1', a: 'thing2', b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(),
                       BSON(""
@@ -102,14 +104,16 @@ TEST(SortKeyGeneratorTest, ExtractStringSortKeyWithCollatorUsesComparisonKey) {
 TEST(SortKeyGeneratorTest, CollatorHasNoEffectWhenExtractingNonStringSortKey) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
-    auto sortKey = sortKeyGen->getSortKey(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr);
+    auto sortKey =
+        sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, z: 10, a: 6, b: 16}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 6));
 }
 
 TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysChoosesCorrectKey) {
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << -1), nullptr);
-    auto sortKey = sortKeyGen->getSortKey(fromjson("{_id: 0, a: [1, 2, 3, 4]}"), nullptr);
+    auto sortKey =
+        sortKeyGen->getSortKeyFromDocument(fromjson("{_id: 0, a: [1, 2, 3, 4]}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 4));
 }
@@ -117,8 +121,8 @@ TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysChoosesCorrectKey) {
 TEST(SortKeyGeneratorTest, EnsureSortKeyGenerationForArraysRespectsCollation) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
-    auto sortKey =
-        sortKeyGen->getSortKey(fromjson("{_id: 0, a: ['aaz', 'zza', 'yya', 'zzb']}"), nullptr);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(
+        fromjson("{_id: 0, a: ['aaz', 'zza', 'yya', 'zzb']}"), nullptr);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(),
                       BSON(""
@@ -127,7 +131,7 @@ TEST(SortKeyGeneratorTest, EnsureSortKeyGenerationForArraysRespectsCollation) {
 
 TEST(SortKeyGeneratorTest, SortKeyGenerationForArraysRespectsCompoundOrdering) {
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a.b" << 1 << "a.c" << -1), nullptr);
-    auto sortKey = sortKeyGen->getSortKey(
+    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());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 0 << "" << 3));
@@ -176,7 +180,7 @@ DEATH_TEST(SortKeyGeneratorTest,
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
                                                                           << "textScore")),
                                                          nullptr);
-    uassertStatusOK(sortKeyGen->getSortKey(BSONObj{}, nullptr).getStatus());
+    uassertStatusOK(sortKeyGen->getSortKeyFromDocument(BSONObj{}, nullptr).getStatus());
 }
 
 DEATH_TEST(SortKeyGeneratorTest,
@@ -185,7 +189,7 @@ DEATH_TEST(SortKeyGeneratorTest,
     auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << BSON("$meta"
                                                                           << "randVal")),
                                                          nullptr);
-    uassertStatusOK(sortKeyGen->getSortKey(BSONObj{}, nullptr).getStatus());
+    uassertStatusOK(sortKeyGen->getSortKeyFromDocument(BSONObj{}, nullptr).getStatus());
 }
 
 TEST(SortKeyGeneratorTest, CanGenerateKeysForTextScoreMetaSort) {
@@ -194,7 +198,7 @@ TEST(SortKeyGeneratorTest, CanGenerateKeysForTextScoreMetaSort) {
                                                          nullptr);
     SortKeyGenerator::Metadata metadata;
     metadata.textScore = 1.5;
-    auto sortKey = sortKeyGen->getSortKey(BSONObj{}, &metadata);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(BSONObj{}, &metadata);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 1.5));
 }
@@ -205,7 +209,7 @@ TEST(SortKeyGeneratorTest, CanGenerateKeysForRandValMetaSort) {
                                                          nullptr);
     SortKeyGenerator::Metadata metadata;
     metadata.randVal = 0.3;
-    auto sortKey = sortKeyGen->getSortKey(BSONObj{}, &metadata);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(BSONObj{}, &metadata);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(), BSON("" << 0.3));
 }
@@ -217,11 +221,102 @@ TEST(SortKeyGeneratorTest, CanGenerateKeysForCompoundMetaSort) {
     SortKeyGenerator::Metadata metadata;
     metadata.randVal = 0.3;
     metadata.textScore = 1.5;
-    auto sortKey = sortKeyGen->getSortKey(BSON("a" << 4 << "d" << 5), &metadata);
+    auto sortKey = sortKeyGen->getSortKeyFromDocument(BSON("a" << 4 << "d" << 5), &metadata);
     ASSERT_OK(sortKey.getStatus());
     ASSERT_BSONOBJ_EQ(sortKey.getValue(),
                       BSON("" << 4 << "" << 0.3 << "" << 1.5 << "" << 5 << "" << 1.5));
 }
 
+// A test fixture which creates a WorkingSet and allocates a WorkingSetMember inside of it. Used for
+// testing sort key generation against a working set member.
+class SortKeyGeneratorWorkingSetTest : public mongo::unittest::Test {
+public:
+    explicit SortKeyGeneratorWorkingSetTest()
+        : _wsid(_workingSet.allocate()), _member(_workingSet.get(_wsid)) {}
+
+    void setRecordIdAndObj(BSONObj obj) {
+        _member->obj = {SnapshotId(), std::move(obj)};
+        _workingSet.transitionToRecordIdAndObj(_wsid);
+    }
+
+    void setOwnedObj(BSONObj obj) {
+        _member->obj = {SnapshotId(), std::move(obj)};
+        _workingSet.transitionToOwnedObj(_wsid);
+    }
+
+    void setRecordIdAndIdx(BSONObj keyPattern, BSONObj key) {
+        _member->keyData.push_back(IndexKeyDatum(std::move(keyPattern), std::move(key), nullptr));
+        _workingSet.transitionToRecordIdAndIdx(_wsid);
+    }
+
+    WorkingSetMember& member() {
+        return *_member;
+    }
+
+private:
+    WorkingSet _workingSet;
+    WorkingSetID _wsid;
+    WorkingSetMember* _member;
+};
+
+TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithObj) {
+    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    setRecordIdAndObj(BSON("x" << 1 << "a" << 2 << "y" << 3));
+    auto sortKey = sortKeyGen->getSortKey(member());
+    ASSERT_OK(sortKey);
+    ASSERT_BSONOBJ_EQ(BSON("" << 2), sortKey.getValue());
+}
+
+TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithOwnedObj) {
+    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3));
+    auto sortKey = sortKeyGen->getSortKey(member());
+    ASSERT_OK(sortKey);
+    ASSERT_BSONOBJ_EQ(BSON("" << 2), sortKey.getValue());
+}
+
+TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateKeyFromWSMForTextScoreMetaSort) {
+    BSONObj pattern = fromjson("{a: 1, b: {$meta: 'textScore'}, c: -1}}");
+    auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr);
+    setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3 << "c" << BSON_ARRAY(4 << 5 << 6)));
+    member().addComputed(new TextScoreComputedData(9.9));
+    auto sortKey = sortKeyGen->getSortKey(member());
+    ASSERT_OK(sortKey);
+    ASSERT_BSONOBJ_EQ(BSON("" << 2 << "" << 9.9 << "" << 6), sortKey.getValue());
+}
+
+TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyState) {
+    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), nullptr);
+    setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << 2 << "" << 3));
+    auto sortKey = sortKeyGen->getSortKey(member());
+    ASSERT_OK(sortKey);
+    ASSERT_BSONOBJ_EQ(BSON("" << 2), sortKey.getValue());
+}
+
+TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyStateWithCollator) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    auto sortKeyGen = std::make_unique<SortKeyGenerator>(BSON("a" << 1), &collator);
+    setRecordIdAndIdx(BSON("a" << 1 << "b" << 1),
+                      BSON(""
+                           << "string1"
+                           << ""
+                           << "string2"));
+    auto sortKey = sortKeyGen->getSortKey(member());
+    ASSERT_OK(sortKey);
+    ASSERT_BSONOBJ_EQ(BSON(""
+                           << "1gnirts"),
+                      sortKey.getValue());
+}
+
+DEATH_TEST_F(SortKeyGeneratorWorkingSetTest,
+             DeathOnAttemptToGetSortKeyFromIndexKeyWithMetadata,
+             "Invariant failure !_sortHasMeta") {
+    BSONObj pattern = fromjson("{z: {$meta: 'textScore'}}");
+    auto sortKeyGen = std::make_unique<SortKeyGenerator>(pattern, nullptr);
+    setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << 2 << "" << 3));
+    member().addComputed(new TextScoreComputedData(9.9));
+    MONGO_COMPILER_VARIABLE_UNUSED auto ignored = sortKeyGen->getSortKey(member());
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_sort.cpp b/src/mongo/db/pipeline/document_source_sort.cpp
index af89701d885..65f3d3bf8a4 100644
--- a/src/mongo/db/pipeline/document_source_sort.cpp
+++ b/src/mongo/db/pipeline/document_source_sort.cpp
@@ -347,7 +347,7 @@ BSONObj DocumentSourceSort::extractKeyWithArray(const Document& doc) const {
     // Convert the Document to a BSONObj, but only do the conversion for the paths we actually need.
     // Then run the result through the SortKeyGenerator to obtain the final sort key.
     auto bsonDoc = _sortExecutor->sortPattern().documentToBsonWithSortPaths(doc);
-    return uassertStatusOK(_sortKeyGen->getSortKey(std::move(bsonDoc), &metadata));
+    return uassertStatusOK(_sortKeyGen->getSortKeyFromDocument(bsonDoc, &metadata));
 }
 
 std::pair<Value, Document> DocumentSourceSort::extractSortKey(Document&& doc) const {