SERVER-42836 Fast path for sort key generation of WorkingSetMembers

(Also includes a comment fixup in document_value_test.cpp that was
intended for a previous commit.)

3 files changed, 98 insertions, 118 deletions

diff --git a/src/mongo/db/index/sort_key_generator.cpp b/src/mongo/db/index/sort_key_generator.cpp
index 97abfe097a7..a09622bb316 100644
--- a/src/mongo/db/index/sort_key_generator.cpp
+++ b/src/mongo/db/index/sort_key_generator.cpp
@@ -76,30 +76,15 @@ SortKeyGenerator::SortKeyGenerator(SortPattern sortPattern, const CollatorInterf
                                                        Ordering::make(_sortSpecWithoutMeta));
 }
 
-// TODO (SERVER-42836): Once WorkingSetMember objects store a Document (SERVER-42181), this function
-// will be able to use the Document overload of computeSortKeyFromDocument, and it will be able to
-// store the text score with the Document instead of in a separate SortKeyGenerator::Metadata
-// object.
-StatusWith<Value> SortKeyGenerator::computeSortKey(const WorkingSetMember& wsm) const {
+Value SortKeyGenerator::computeSortKey(const WorkingSetMember& wsm) const {
     if (wsm.hasObj()) {
-        SortKeyGenerator::Metadata metadata;
-        if (_sortHasMeta && wsm.metadata().hasTextScore()) {
-            metadata.textScore = wsm.metadata().getTextScore();
-        }
-        auto statusWithSortKeyObj = computeSortKeyFromDocument(wsm.doc.value().toBson(), &metadata);
-        if (!statusWithSortKeyObj.isOK()) {
-            return statusWithSortKeyObj.getStatus();
-        }
-
-        return DocumentMetadataFields::deserializeSortKey(isSingleElementKey(),
-                                                          statusWithSortKeyObj.getValue());
+        return computeSortKeyFromDocument(wsm.doc.value(), wsm.metadata());
     }
 
     return computeSortKeyFromIndexKey(wsm);
 }
 
-StatusWith<Value> SortKeyGenerator::computeSortKeyFromIndexKey(
-    const WorkingSetMember& member) const {
+Value SortKeyGenerator::computeSortKeyFromIndexKey(const WorkingSetMember& member) const {
     invariant(member.getState() == WorkingSetMember::RID_AND_IDX);
     invariant(!_sortHasMeta);
 
@@ -118,16 +103,9 @@ StatusWith<Value> SortKeyGenerator::computeSortKeyFromIndexKey(
     return DocumentMetadataFields::deserializeSortKey(isSingleElementKey(), objBuilder.obj());
 }
 
-StatusWith<BSONObj> SortKeyGenerator::computeSortKeyFromDocument(const BSONObj& obj,
-                                                                 const Metadata* metadata) const {
-    if (_sortHasMeta) {
-        invariant(metadata);
-    }
-
-    auto sortKeyNoMetadata = computeSortKeyFromDocumentWithoutMetadata(obj);
-    if (!sortKeyNoMetadata.isOK()) {
-        return sortKeyNoMetadata;
-    }
+BSONObj SortKeyGenerator::computeSortKeyFromDocument(const BSONObj& obj,
+                                                     const DocumentMetadataFields& metadata) const {
+    auto sortKeyNoMetadata = uassertStatusOK(computeSortKeyFromDocumentWithoutMetadata(obj));
 
     if (!_sortHasMeta) {
         // We don't have to worry about $meta sort, so the index key becomes the sort key.
@@ -137,24 +115,27 @@ StatusWith<BSONObj> SortKeyGenerator::computeSortKeyFromDocument(const BSONObj&
     BSONObjBuilder mergedKeyBob;
 
     // Merge metadata into the key.
-    BSONObjIterator sortKeyIt(sortKeyNoMetadata.getValue());
+    BSONObjIterator sortKeyIt(sortKeyNoMetadata);
     for (auto& part : _sortPattern) {
         if (part.fieldPath) {
             invariant(sortKeyIt.more());
             mergedKeyBob.append(sortKeyIt.next());
             continue;
         }
+
+        // Create a Document that represents the input object and its metadata together, so we can
+        // use it to evaluate the ExpressionMeta for this part of the sort pattern. This operation
+        // copies the data in 'metadata' but not any of the data in the 'obj' BSON.
+        MutableDocument documentWithMetdata(Document{obj});
+        documentWithMetdata.setMetadata(DocumentMetadataFields(metadata));
+
         invariant(part.expression);
-        switch (part.expression->getMetaType()) {
-            case DocumentMetadataFields::MetaType::kTextScore: {
-                mergedKeyBob.append("", metadata->textScore);
-                continue;
-            }
-            case DocumentMetadataFields::MetaType::kRandVal: {
-                mergedKeyBob.append("", metadata->randVal);
-                continue;
-            }
-            default: { MONGO_UNREACHABLE; }
+        auto value =
+            part.expression->evaluate(documentWithMetdata.freeze(), nullptr /* variables */);
+        if (!value.missing()) {
+            value.addToBsonObj(&mergedKeyBob, ""_sd);
+        } else {
+            mergedKeyBob.appendNull("");
         }
     }
 
@@ -231,7 +212,9 @@ Value SortKeyGenerator::getCollationComparisonKey(const Value& val) const {
 }
 
 StatusWith<Value> SortKeyGenerator::extractKeyPart(
-    const Document& doc, const SortPattern::SortPatternPart& patternPart) const {
+    const Document& doc,
+    const DocumentMetadataFields& metadata,
+    const SortPattern::SortPatternPart& patternPart) const {
     Value plainKey;
     if (patternPart.fieldPath) {
         invariant(!patternPart.expression);
@@ -243,22 +226,28 @@ StatusWith<Value> SortKeyGenerator::extractKeyPart(
         plainKey = key.getValue();
     } else {
         invariant(patternPart.expression);
+        // ExpressionMeta expects metadata to be attached to the document.
+        MutableDocument documentWithMetadata(doc);
+        documentWithMetadata.setMetadata(DocumentMetadataFields(metadata));
+
         // ExpressionMeta does not use Variables.
-        plainKey = patternPart.expression->evaluate(doc, nullptr /* variables */);
+        plainKey = patternPart.expression->evaluate(documentWithMetadata.freeze(),
+                                                    nullptr /* variables */);
     }
 
-    return getCollationComparisonKey(plainKey);
+    return plainKey.missing() ? Value{BSONNULL} : getCollationComparisonKey(plainKey);
 }
 
-StatusWith<Value> SortKeyGenerator::extractKeyFast(const Document& doc) const {
+StatusWith<Value> SortKeyGenerator::extractKeyFast(const Document& doc,
+                                                   const DocumentMetadataFields& metadata) const {
     if (_sortPattern.isSingleElementKey()) {
-        return extractKeyPart(doc, _sortPattern[0]);
+        return extractKeyPart(doc, metadata, _sortPattern[0]);
     }
 
     std::vector<Value> keys;
     keys.reserve(_sortPattern.size());
     for (auto&& keyPart : _sortPattern) {
-        auto extractedKey = extractKeyPart(doc, keyPart);
+        auto extractedKey = extractKeyPart(doc, metadata, keyPart);
         if (!extractedKey.isOK()) {
             // We can't use the fast path, so bail out.
             return extractedKey;
@@ -269,24 +258,18 @@ StatusWith<Value> SortKeyGenerator::extractKeyFast(const Document& doc) const {
     return Value{std::move(keys)};
 }
 
-BSONObj SortKeyGenerator::extractKeyWithArray(const Document& doc) const {
-    SortKeyGenerator::Metadata metadata;
-    if (doc.metadata().hasTextScore()) {
-        metadata.textScore = doc.metadata().getTextScore();
-    }
-    if (doc.metadata().hasRandVal()) {
-        metadata.randVal = doc.metadata().getRandVal();
-    }
-
+BSONObj SortKeyGenerator::extractKeyWithArray(const Document& doc,
+                                              const DocumentMetadataFields& metadata) 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 = _sortPattern.documentToBsonWithSortPaths(doc);
-    return uassertStatusOK(computeSortKeyFromDocument(bsonDoc, &metadata));
+    return computeSortKeyFromDocument(bsonDoc, metadata);
 }
 
-Value SortKeyGenerator::computeSortKeyFromDocument(const Document& doc) const {
+Value SortKeyGenerator::computeSortKeyFromDocument(const Document& doc,
+                                                   const DocumentMetadataFields& metadata) const {
     // This fast pass directly generates a Value.
-    auto fastKey = extractKeyFast(doc);
+    auto fastKey = extractKeyFast(doc, metadata);
     if (fastKey.isOK()) {
         return std::move(fastKey.getValue());
     }
@@ -295,7 +278,7 @@ Value SortKeyGenerator::computeSortKeyFromDocument(const Document& doc) const {
     // form like BSONObj {'': 1, '': [2, 3]}) and converts it to a Value (Value [1, [2, 3]] in the
     // earlier example).
     return DocumentMetadataFields::deserializeSortKey(_sortPattern.isSingleElementKey(),
-                                                      extractKeyWithArray(doc));
+                                                      extractKeyWithArray(doc, metadata));
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/index/sort_key_generator.h b/src/mongo/db/index/sort_key_generator.h
index b6ea016139a..dd4b5031d05 100644
--- a/src/mongo/db/index/sort_key_generator.h
+++ b/src/mongo/db/index/sort_key_generator.h
@@ -41,16 +41,6 @@ namespace mongo {
 class SortKeyGenerator {
 public:
     /**
-     * Metadata about a document which is needed to produce keys for $meta sort. The client of the
-     * SortKeyGenerator must provide this metadata in order to correctly obtain sort keys for sort
-     * patterns with $meta.
-     */
-    struct Metadata {
-        double textScore = 0.0;
-        double randVal = 0.0;
-    };
-
-    /**
      * 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.
@@ -58,75 +48,87 @@ public:
     SortKeyGenerator(SortPattern sortPattern, 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.
+     * Returns the key which should be used to sort the WorkingSetMember or throws 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<Value> computeSortKey(const WorkingSetMember&) const;
+    Value computeSortKey(const WorkingSetMember&) const;
 
     /**
-     * Returns the sort key for the input 'doc' as a Value. When the sort pattern has multiple
-     * components, the resulting sort key is a an Array-typed Value with one element for each
-     * component. For sort pattern with just one component, the sort key is a Value that represents
-     * the single element to sort on (which may or may not itself be an array).
+     * Returns the sort key for the input 'doc' as a Value or throws if no key could be generated.
+     * When the sort pattern has multiple components, the resulting sort key is an Array-typed Value
+     * with one element for each component. For sort patterns with just one component, the sort key
+     * is a Value that represents the single element to sort on (which may or may not itself be an
+     * array).
      *
      * The sort key is computed based on the sort pattern, the contents of the document, and if
-     * required by $meta sort specifiers, metadata in the Document. This function throws if it
-     * cannot compute the sort pattern.
+     * required by $meta sort specifiers, metadata in the Document.
      */
-    Value computeSortKeyFromDocument(const Document& doc) const;
+    Value computeSortKeyFromDocument(const Document& doc) const {
+        return computeSortKeyFromDocument(doc, doc.metadata());
+    }
 
     bool isSingleElementKey() const {
         return _sortPattern.isSingleElementKey();
     }
 
 private:
-    /**
-     * Returns the key which should be used to sort 'obj', or a non-OK status if no key could be
-     * generated.
-     *
-     * The caller must supply the appropriate 'metadata' in the case that the sort pattern includes
-     * a $meta sort (i.e. if sortHasMeta() is true). These values are filled in at the corresponding
-     * positions in the sort key.
-     */
-    StatusWith<BSONObj> computeSortKeyFromDocument(const BSONObj& obj, const Metadata*) const;
+    // Returns the sort key for the input 'doc' as a Value.
+    //
+    // Note that this function will ignore any metadata (e.g., textScore, randVal), in 'doc' but
+    // will instead read from the 'metadata' variable. When the metadata is contained in the 'doc'
+    // input, callers can use the public overload of this function.
+    Value computeSortKeyFromDocument(const Document& doc,
+                                     const DocumentMetadataFields& metadata) const;
+
+    // Returns the key which should be used to sort 'obj' or throws an exception if no key could be
+    // generated.
+    //
+    // The caller must supply the appropriate 'metadata' in the case that the sort pattern includes
+    // a $meta sort (i.e. if sortHasMeta() is true). These values are filled in at the corresponding
+    // positions in the sort key.
+    BSONObj computeSortKeyFromDocument(const BSONObj& obj,
+                                       const DocumentMetadataFields& metadata) 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<Value> computeSortKeyFromIndexKey(const WorkingSetMember& member) const;
+    Value computeSortKeyFromIndexKey(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> computeSortKeyFromDocumentWithoutMetadata(const BSONObj& obj) const;
 
-    /**
-     * Returns the sort key for 'doc' based on the SortPattern, or ErrorCodes::InternalError if an
-     * array is encountered during sort key generation.
-     */
-    StatusWith<Value> extractKeyFast(const Document& doc) const;
-
-    /**
-     * Extracts the sort key component described by 'keyPart' from 'doc' and returns it. Returns
-     * ErrorCodes::InternalError if the path for 'keyPart' contains an array in 'doc'.
-     */
+    // Returns the sort key for 'doc' based on the SortPattern, or ErrorCodes::InternalError if an
+    // array is encountered during sort key generation.
+    //
+    // Note that this function will ignore any metadata (e.g., textScore, randVal), in 'doc' but
+    // will instead read from the 'metadata' variable.
+    StatusWith<Value> extractKeyFast(const Document& doc,
+                                     const DocumentMetadataFields& metadata) const;
+
+    // Extracts the sort key component described by 'keyPart' from 'doc' and returns it. Returns
+    // ErrorCodes::InternalError if the path for 'keyPart' contains an array in 'doc'.
+    //
+    // Note that this function will ignore any metadata (e.g., textScore, randVal), in 'doc' but
+    // will instead read from the 'metadata' variable.
     StatusWith<Value> extractKeyPart(const Document& doc,
+                                     const DocumentMetadataFields& metadata,
                                      const SortPattern::SortPatternPart& keyPart) const;
 
-    /**
-     * Returns the sort key for 'doc' based on the SortPattern. Note this is in the BSONObj format -
-     * with empty field names.
-     */
-    BSONObj extractKeyWithArray(const Document& doc) const;
+    // Returns the sort key for 'doc' based on the SortPattern. Note this is in the BSONObj format -
+    // with empty field names.
+    //
+    // Note that this function will ignore any metadata (e.g., textScore, randVal), in 'doc' but
+    // will instead read from the 'metadata' variable.
+    BSONObj extractKeyWithArray(const Document& doc, const DocumentMetadataFields& metadata) const;
 
-    /**
-     * Returns the comparison key used to sort 'val' with collation. Note that these comparison keys
-     * should always be sorted with the simple (i.e. binary) collation.
-     */
+    // Returns the comparison key used to sort 'val' with collation. Note that these comparison keys
+    // should always be sorted with the simple (i.e. binary) collation.
     Value getCollationComparisonKey(const Value& val) 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 0e6681926dc..29ca2039153 100644
--- a/src/mongo/db/index/sort_key_generator_test.cpp
+++ b/src/mongo/db/index/sort_key_generator_test.cpp
@@ -266,16 +266,14 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithObj)
     auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     setRecordIdAndObj(BSON("x" << 1 << "a" << 2 << "y" << 3));
     auto sortKey = sortKeyGen->computeSortKey(member());
-    ASSERT_OK(sortKey);
-    ASSERT_VALUE_EQ(Value(2), sortKey.getValue());
+    ASSERT_VALUE_EQ(Value(2), sortKey);
 }
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGetSortKeyFromWorkingSetMemberWithOwnedObj) {
     auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3));
     auto sortKey = sortKeyGen->computeSortKey(member());
-    ASSERT_OK(sortKey);
-    ASSERT_VALUE_EQ(Value(2), sortKey.getValue());
+    ASSERT_VALUE_EQ(Value(2), sortKey);
 }
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateKeyFromWSMForTextScoreMetaSort) {
@@ -284,16 +282,14 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateKeyFromWSMForTextScoreMetaSort
     setOwnedObj(BSON("x" << 1 << "a" << 2 << "y" << 3 << "c" << BSON_ARRAY(4 << 5 << 6)));
     member().metadata().setTextScore(9.9);
     auto sortKey = sortKeyGen->computeSortKey(member());
-    ASSERT_OK(sortKey);
-    ASSERT_VALUE_EQ(Value({Value(2), Value(9.9), Value(6)}), sortKey.getValue());
+    ASSERT_VALUE_EQ(Value({Value(2), Value(9.9), Value(6)}), sortKey);
 }
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyState) {
     auto sortKeyGen = makeSortKeyGen(BSON("a" << 1), nullptr);
     setRecordIdAndIdx(BSON("a" << 1 << "b" << 1), BSON("" << 2 << "" << 3));
     auto sortKey = sortKeyGen->computeSortKey(member());
-    ASSERT_OK(sortKey);
-    ASSERT_VALUE_EQ(Value(2), sortKey.getValue());
+    ASSERT_VALUE_EQ(Value(2), sortKey);
 }
 
 TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyStateWithCollator) {
@@ -305,8 +301,7 @@ TEST_F(SortKeyGeneratorWorkingSetTest, CanGenerateSortKeyFromWSMInIndexKeyStateW
                            << ""
                            << "string2"));
     auto sortKey = sortKeyGen->computeSortKey(member());
-    ASSERT_OK(sortKey);
-    ASSERT_VALUE_EQ(Value("1gnirts"_sd), sortKey.getValue());
+    ASSERT_VALUE_EQ(Value("1gnirts"_sd), sortKey);
 }
 
 DEATH_TEST_F(SortKeyGeneratorWorkingSetTest,