SERVER-40016 Add searchScore metadata handling

5 files changed, 53 insertions, 2 deletions

diff --git a/src/mongo/db/pipeline/document.cpp b/src/mongo/db/pipeline/document.cpp
index 7d0d23ed24c..e7d74aadbfe 100644
--- a/src/mongo/db/pipeline/document.cpp
+++ b/src/mongo/db/pipeline/document.cpp
@@ -49,7 +49,8 @@ const std::vector<StringData> Document::allMetadataFieldNames = {Document::metaF
                                                                  Document::metaFieldRandVal,
                                                                  Document::metaFieldSortKey,
                                                                  Document::metaFieldGeoNearDistance,
-                                                                 Document::metaFieldGeoNearPoint};
+                                                                 Document::metaFieldGeoNearPoint,
+                                                                 Document::metaFieldSearchScore};
 
 Position DocumentStorage::findField(StringData requested) const {
     int reqSize = requested.size();  // get size calculation out of the way if needed
@@ -224,6 +225,7 @@ intrusive_ptr<DocumentStorage> DocumentStorage::clone() const {
         out->_sortKey = _sortKey.getOwned();
         out->_geoNearDistance = _geoNearDistance;
         out->_geoNearPoint = _geoNearPoint.getOwned();
+        out->_searchScore = _searchScore;
     }
 
     return out;
@@ -289,6 +291,7 @@ constexpr StringData Document::metaFieldRandVal;
 constexpr StringData Document::metaFieldSortKey;
 constexpr StringData Document::metaFieldGeoNearDistance;
 constexpr StringData Document::metaFieldGeoNearPoint;
+constexpr StringData Document::metaFieldSearchScore;
 
 BSONObj Document::toBsonWithMetaData() const {
     BSONObjBuilder bb;
@@ -303,6 +306,8 @@ BSONObj Document::toBsonWithMetaData() const {
         bb.append(metaFieldGeoNearDistance, getGeoNearDistance());
     if (hasGeoNearPoint())
         getGeoNearPoint().addToBsonObj(&bb, metaFieldGeoNearPoint);
+    if (hasSearchScore())
+        bb.append(metaFieldSearchScore, getSearchScore());
     return bb.obj();
 }
 
@@ -317,6 +322,9 @@ Document Document::fromBsonWithMetaData(const BSONObj& bson) {
             if (fieldName == metaFieldTextScore) {
                 md.setTextScore(elem.Double());
                 continue;
+            } else if (fieldName == metaFieldSearchScore) {
+                md.setSearchScore(elem.Double());
+                continue;
             } else if (fieldName == metaFieldRandVal) {
                 md.setRandMetaField(elem.Double());
                 continue;
@@ -514,6 +522,10 @@ void Document::serializeForSorter(BufBuilder& buf) const {
         buf.appendNum(char(DocumentStorage::MetaType::SORT_KEY + 1));
         getSortKeyMetaField().appendSelfToBufBuilder(buf);
     }
+    if (hasSearchScore()) {
+        buf.appendNum(char(DocumentStorage::MetaType::SEARCH_SCORE + 1));
+        buf.appendNum(getSearchScore());
+    }
     buf.appendNum(char(0));
 }
 
@@ -533,6 +545,8 @@ Document Document::deserializeForSorter(BufReader& buf, const SorterDeserializeS
         } else if (marker == char(DocumentStorage::MetaType::SORT_KEY) + 1) {
             doc.setSortKeyMetaField(
                 BSONObj::deserializeForSorter(buf, BSONObj::SorterDeserializeSettings()));
+        } else if (marker == char(DocumentStorage::MetaType::SEARCH_SCORE) + 1) {
+            doc.setSearchScore(buf.read<LittleEndian<double>>());
         } else {
             uasserted(28744, "Unrecognized marker, unable to deserialize buffer");
         }
diff --git a/src/mongo/db/pipeline/document.h b/src/mongo/db/pipeline/document.h
index 0c64c108329..6d765a24c65 100644
--- a/src/mongo/db/pipeline/document.h
+++ b/src/mongo/db/pipeline/document.h
@@ -96,6 +96,7 @@ public:
     static constexpr StringData metaFieldSortKey = "$sortKey"_sd;
     static constexpr StringData metaFieldGeoNearDistance = "$dis"_sd;
     static constexpr StringData metaFieldGeoNearPoint = "$pt"_sd;
+    static constexpr StringData metaFieldSearchScore = "$searchScore"_sd;
 
     static const std::vector<StringData> allMetadataFieldNames;
 
@@ -283,6 +284,13 @@ public:
         return storage().getGeoNearPoint();
     }
 
+    bool hasSearchScore() const {
+        return storage().hasSearchScore();
+    }
+    double getSearchScore() const {
+        return storage().getSearchScore();
+    }
+
     /// members for Sorter
     struct SorterDeserializeSettings {};  // unused
     void serializeForSorter(BufBuilder& buf) const;
@@ -542,6 +550,10 @@ public:
         storage().setGeoNearPoint(std::move(point));
     }
 
+    void setSearchScore(double score) {
+        storage().setSearchScore(score);
+    }
+
     /** Convert to a read-only document and release reference.
      *
      *  Call this to indicate that you are done with this Document and will
diff --git a/src/mongo/db/pipeline/document_internal.h b/src/mongo/db/pipeline/document_internal.h
index b5030b77693..71b28d2aeb8 100644
--- a/src/mongo/db/pipeline/document_internal.h
+++ b/src/mongo/db/pipeline/document_internal.h
@@ -192,7 +192,8 @@ public:
           _metaFields(),
           _textScore(0),
           _randVal(0),
-          _geoNearDistance(0) {}
+          _geoNearDistance(0),
+          _searchScore(0) {}
 
     ~DocumentStorage();
 
@@ -202,6 +203,7 @@ public:
         SORT_KEY,
         GEONEAR_DIST,
         GEONEAR_POINT,
+        SEARCH_SCORE,
 
         // New fields must be added before the NUM_FIELDS sentinel.
         NUM_FIELDS
@@ -296,6 +298,9 @@ public:
         if (source.hasGeoNearPoint()) {
             setGeoNearPoint(source.getGeoNearPoint());
         }
+        if (source.hasSearchScore()) {
+            setSearchScore(source.getSearchScore());
+        }
     }
 
     bool hasTextScore() const {
@@ -353,6 +358,17 @@ public:
         _geoNearPoint = std::move(point);
     }
 
+    bool hasSearchScore() const {
+        return _metaFields.test(MetaType::SEARCH_SCORE);
+    }
+    double getSearchScore() const {
+        return _searchScore;
+    }
+    void setSearchScore(double score) {
+        _metaFields.set(MetaType::SEARCH_SCORE);
+        _searchScore = score;
+    }
+
 private:
     /// Same as lastElement->next() or firstElement() if empty.
     const ValueElement* end() const {
@@ -437,6 +453,7 @@ private:
     BSONObj _sortKey;
     double _geoNearDistance;
     Value _geoNearPoint;
+    double _searchScore;
     // When adding a field, make sure to update clone() method
 
     // Defined in document.cpp
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index d8ab5bce397..87a2d841e5e 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -2517,6 +2517,8 @@ intrusive_ptr<Expression> ExpressionMeta::parse(
         return new ExpressionMeta(expCtx, MetaType::TEXT_SCORE);
     } else if (expr.valueStringData() == "randVal") {
         return new ExpressionMeta(expCtx, MetaType::RAND_VAL);
+    } else if (expr.valueStringData() == "searchScore") {
+        return new ExpressionMeta(expCtx, MetaType::SEARCH_SCORE);
     } else {
         uasserted(17308, "Unsupported argument to $meta: " + expr.String());
     }
@@ -2534,6 +2536,9 @@ Value ExpressionMeta::serialize(bool explain) const {
         case MetaType::RAND_VAL:
             return Value(DOC("$meta"
                              << "randVal"_sd));
+        case MetaType::SEARCH_SCORE:
+            return Value(DOC("$meta"
+                             << "searchScore"_sd));
     }
     MONGO_UNREACHABLE;
 }
@@ -2544,6 +2549,8 @@ Value ExpressionMeta::evaluate(const Document& root) const {
             return root.hasTextScore() ? Value(root.getTextScore()) : Value();
         case MetaType::RAND_VAL:
             return root.hasRandMetaField() ? Value(root.getRandMetaField()) : Value();
+        case MetaType::SEARCH_SCORE:
+            return root.hasSearchScore() ? Value(root.getSearchScore()) : Value();
     }
     MONGO_UNREACHABLE;
 }
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 47829732ee4..ee468b550c1 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1631,6 +1631,7 @@ private:
     enum MetaType {
         TEXT_SCORE,
         RAND_VAL,
+        SEARCH_SCORE,
     };
 
     ExpressionMeta(const boost::intrusive_ptr<ExpressionContext>& expCtx, MetaType metaType);