SERVER-40968 Core DocumentStorage changes for Document/Value optimization

4 files changed, 10 insertions, 15 deletions

diff --git a/src/mongo/s/query/cluster_client_cursor_impl.cpp b/src/mongo/s/query/cluster_client_cursor_impl.cpp
index 3db2225e1ae..ecfb31a2a46 100644
--- a/src/mongo/s/query/cluster_client_cursor_impl.cpp
+++ b/src/mongo/s/query/cluster_client_cursor_impl.cpp
@@ -242,7 +242,7 @@ std::unique_ptr<RouterExecStage> ClusterClientCursorImpl::buildMergerPlan(
     if (hasSort) {
         // Strip out the sort key after sorting.
         root = std::make_unique<RouterStageRemoveMetadataFields>(
-            opCtx, std::move(root), std::vector<StringData>{AsyncResultsMerger::kSortKeyField});
+            opCtx, std::move(root), StringDataSet{AsyncResultsMerger::kSortKeyField});
     }
 
     return root;
diff --git a/src/mongo/s/query/router_stage_remove_metadata_fields.cpp b/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
index 70824081b50..c091a083c64 100644
--- a/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
+++ b/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
@@ -39,9 +39,7 @@
 namespace mongo {
 
 RouterStageRemoveMetadataFields::RouterStageRemoveMetadataFields(
-    OperationContext* opCtx,
-    std::unique_ptr<RouterExecStage> child,
-    std::vector<StringData> metadataFields)
+    OperationContext* opCtx, std::unique_ptr<RouterExecStage> child, StringDataSet metadataFields)
     : RouterExecStage(opCtx, std::move(child)), _metaFields(std::move(metadataFields)) {
     for (auto&& fieldName : _metaFields) {
         invariant(fieldName[0] == '$');  // We use this information to optimize next().
@@ -57,9 +55,7 @@ StatusWith<ClusterQueryResult> RouterStageRemoveMetadataFields::next(
 
     BSONObjIterator iterator(*childResult.getValue().getResult());
     // Find the first field that we need to remove.
-    while (iterator.more() && (*iterator).fieldName()[0] != '$' &&
-           std::find(_metaFields.begin(), _metaFields.end(), (*iterator).fieldNameStringData()) ==
-               _metaFields.end()) {
+    while (iterator.more() && (*iterator).fieldName()[0] != '$') {
         ++iterator;
     }
 
@@ -80,8 +76,7 @@ StatusWith<ClusterQueryResult> RouterStageRemoveMetadataFields::next(
     // Copy any remaining fields that are not metadata. We expect metadata fields are likely to be
     // at the end of the document, so there is likely nothing else to copy.
     while ((++iterator).more()) {
-        if (std::find(_metaFields.begin(), _metaFields.end(), (*iterator).fieldNameStringData()) ==
-            _metaFields.end()) {
+        if (!_metaFields.contains((*iterator).fieldNameStringData())) {
             builder.append(*iterator);
         }
     }
diff --git a/src/mongo/s/query/router_stage_remove_metadata_fields.h b/src/mongo/s/query/router_stage_remove_metadata_fields.h
index 921fe5d4f92..4f49537b2f9 100644
--- a/src/mongo/s/query/router_stage_remove_metadata_fields.h
+++ b/src/mongo/s/query/router_stage_remove_metadata_fields.h
@@ -42,14 +42,14 @@ class RouterStageRemoveMetadataFields final : public RouterExecStage {
 public:
     RouterStageRemoveMetadataFields(OperationContext* opCtx,
                                     std::unique_ptr<RouterExecStage> child,
-                                    std::vector<StringData> fieldsToRemove);
+                                    StringDataSet fieldsToRemove);
 
     StatusWith<ClusterQueryResult> next(ExecContext) final;
 
 private:
     // Use a StringMap so we can look up by StringData - avoiding a string allocation on each field
     // in each object. The value here is meaningless.
-    std::vector<StringData> _metaFields;
+    StringDataSet _metaFields;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp b/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp
index abedcbb6122..a1a7404b486 100644
--- a/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp
+++ b/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp
@@ -110,7 +110,7 @@ TEST(RouterStageRemoveMetadataFieldsTest, PropagatesError) {
     mockStage->queueError(Status(ErrorCodes::BadValue, "bad thing happened"));
 
     auto sortKeyStage = std::make_unique<RouterStageRemoveMetadataFields>(
-        opCtx, std::move(mockStage), std::vector<StringData>{"$sortKey"_sd});
+        opCtx, std::move(mockStage), StringDataSet{"$sortKey"_sd});
 
     auto firstResult = sortKeyStage->next(RouterExecStage::ExecContext::kInitialFind);
     ASSERT_OK(firstResult.getStatus());
@@ -130,7 +130,7 @@ TEST(RouterStageRemoveMetadataFieldsTest, ToleratesMidStreamEOF) {
     mockStage->queueResult(BSON("a" << 2 << "$sortKey" << 1 << "b" << 2));
 
     auto sortKeyStage = std::make_unique<RouterStageRemoveMetadataFields>(
-        opCtx, std::move(mockStage), std::vector<StringData>{"$sortKey"_sd});
+        opCtx, std::move(mockStage), StringDataSet{"$sortKey"_sd});
 
     auto firstResult = sortKeyStage->next(RouterExecStage::ExecContext::kInitialFind);
     ASSERT_OK(firstResult.getStatus());
@@ -158,7 +158,7 @@ TEST(RouterStageRemoveMetadataFieldsTest, RemotesExhausted) {
     mockStage->markRemotesExhausted();
 
     auto sortKeyStage = std::make_unique<RouterStageRemoveMetadataFields>(
-        opCtx, std::move(mockStage), std::vector<StringData>{"$sortKey"_sd});
+        opCtx, std::move(mockStage), StringDataSet{"$sortKey"_sd});
     ASSERT_TRUE(sortKeyStage->remotesExhausted());
 
     auto firstResult = sortKeyStage->next(RouterExecStage::ExecContext::kInitialFind);
@@ -185,7 +185,7 @@ TEST(RouterStageRemoveMetadataFieldsTest, ForwardsAwaitDataTimeout) {
     ASSERT_NOT_OK(mockStage->getAwaitDataTimeout().getStatus());
 
     auto sortKeyStage = std::make_unique<RouterStageRemoveMetadataFields>(
-        opCtx, std::move(mockStage), std::vector<StringData>{"$sortKey"_sd});
+        opCtx, std::move(mockStage), StringDataSet{"$sortKey"_sd});
     ASSERT_OK(sortKeyStage->setAwaitDataTimeout(Milliseconds(789)));
 
     auto awaitDataTimeout = mockStagePtr->getAwaitDataTimeout();