SERVER-38077 Optimize ProjectionNode output for inclusion projections

3 files changed, 22 insertions, 7 deletions

diff --git a/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp b/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp
index 22b5e8df5a8..3be7c403e1c 100644
--- a/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp
+++ b/src/mongo/db/pipeline/parsed_aggregation_projection_node.cpp
@@ -100,22 +100,25 @@ void ProjectionNode::applyProjections(const Document& inputDoc, MutableDocument*
         auto fieldPair = it.next();
         auto fieldName = fieldPair.first.toString();
         if (_projectedFields.count(fieldName)) {
-            outputDoc->setField(fieldName, applyLeafProjectionToValue(fieldPair.second));
+            outputProjectedField(
+                fieldName, applyLeafProjectionToValue(fieldPair.second), outputDoc);
             continue;
         }
 
         auto childIt = _children.find(fieldName);
         if (childIt != _children.end()) {
-            outputDoc->setField(fieldName,
-                                childIt->second->applyProjectionsToValue(fieldPair.second));
+            outputProjectedField(
+                fieldName, childIt->second->applyProjectionsToValue(fieldPair.second), outputDoc);
         }
     }
 
     // Ensure we project all specified fields, including those not present in the input document.
-    const bool shouldProjectNonExistentFields = applyLeafProjectionToValue(Value(true)).missing();
-    for (auto&& fieldName : _projectedFields) {
-        if (shouldProjectNonExistentFields && inputDoc[fieldName].missing()) {
-            outputDoc->setField(fieldName, applyLeafProjectionToValue(inputDoc[fieldName]));
+    // TODO SERVER-37791: This block is only necessary due to a bug in exclusion semantics.
+    if (applyLeafProjectionToValue(Value(true)).missing()) {
+        for (auto&& fieldName : _projectedFields) {
+            if (inputDoc[fieldName].missing()) {
+                outputProjectedField(fieldName, Value(), outputDoc);
+            }
         }
     }
 }
@@ -143,6 +146,10 @@ Value ProjectionNode::applyProjectionsToValue(Value inputValue) const {
     }
 }
 
+void ProjectionNode::outputProjectedField(StringData field, Value val, MutableDocument* doc) const {
+    doc->setField(field, val);
+}
+
 void ProjectionNode::applyExpressions(const Document& root, MutableDocument* outputDoc) const {
     for (auto&& field : _orderToProcessAdditionsAndChildren) {
         auto childIt = _children.find(field);
diff --git a/src/mongo/db/pipeline/parsed_aggregation_projection_node.h b/src/mongo/db/pipeline/parsed_aggregation_projection_node.h
index 517dfe8eaee..8671a289695 100644
--- a/src/mongo/db/pipeline/parsed_aggregation_projection_node.h
+++ b/src/mongo/db/pipeline/parsed_aggregation_projection_node.h
@@ -138,6 +138,9 @@ protected:
     // the output document. Depending on the projection this will be either the value, or "missing".
     virtual Value transformSkippedValueForOutput(const Value&) const = 0;
 
+    // Writes the given value to the output doc, replacing the existing value of 'field' if present.
+    virtual void outputProjectedField(StringData field, Value val, MutableDocument* outDoc) const;
+
     // TODO use StringMap once SERVER-23700 is resolved.
     stdx::unordered_map<std::string, std::unique_ptr<ProjectionNode>> _children;
     stdx::unordered_map<size_t, std::unique_ptr<ProjectionNode>> _arrayBranches;
diff --git a/src/mongo/db/pipeline/parsed_inclusion_projection.h b/src/mongo/db/pipeline/parsed_inclusion_projection.h
index bd11c3e638c..b2a98ae6967 100644
--- a/src/mongo/db/pipeline/parsed_inclusion_projection.h
+++ b/src/mongo/db/pipeline/parsed_inclusion_projection.h
@@ -62,6 +62,11 @@ public:
     void reportDependencies(DepsTracker* deps) const final;
 
 protected:
+    // For inclusions, we can apply an optimization here by simply appending to the output document
+    // via MutableDocument::addField, rather than always checking for existing fields via setField.
+    void outputProjectedField(StringData field, Value val, MutableDocument* outputDoc) const final {
+        outputDoc->addField(field, val);
+    }
     std::unique_ptr<ProjectionNode> makeChild(std::string fieldName) const final {
         return std::make_unique<InclusionNode>(
             _policies, FieldPath::getFullyQualifiedPath(_pathToNode, fieldName));