SERVER-48684 Pipeline stage $set fails to set a non-existent dotted field to an object

(cherry picked from commit 559406f9d52e96f3c10ac3be16c5890714adb2af)

9 files changed, 134 insertions, 22 deletions

diff --git a/jstests/aggregation/sources/addFields/dotted_paths.js b/jstests/aggregation/sources/addFields/dotted_paths.js
new file mode 100644
index 00000000000..e8ecdc39d46
--- /dev/null
+++ b/jstests/aggregation/sources/addFields/dotted_paths.js
@@ -0,0 +1,98 @@
+/**
+ * Test the behavior of $addFields in the presence of a dotted field path.
+ */
+(function() {
+const coll = db.add_fields_dotted_paths;
+coll.drop();
+
+let initObj = {
+    _id: 1,
+    arrayField: [1, {subField: [2, {}]}, [1]],
+    objField: {p: {q: 1}, subArr: [1]},
+    otherField: "value"
+};
+assert.commandWorked(coll.insert(initObj));
+
+function assertAddFieldsResult(projection, expectedResults) {
+    assert.eq(coll.aggregate([{$addFields: projection}]).toArray(), [expectedResults]);
+}
+
+// Test that the value gets overwritten when a field exists at a given path.
+initObj["objField"]["subArr"] = "newValue";
+assertAddFieldsResult({"objField.subArr": "newValue"}, initObj);
+assertAddFieldsResult({"objField.subArr": {$literal: "newValue"}}, initObj);
+
+// Test that a new sub-object is created when a field does not exist at a given path. All the
+// existing sibling fields are retained.
+initObj["objField"] = {
+    p: {q: 1},
+    subArr: [1],  // Existing fields are retained.
+    newSubPath: {b: "newValue"}
+};
+assertAddFieldsResult({"objField.newSubPath.b": {$literal: "newValue"}}, initObj);
+assertAddFieldsResult({"objField.newSubPath.b": "newValue"}, initObj);
+
+// When the value is a nested object.
+const valueWithNestedObject = {
+    newSubObj: [{p: "newValue"}]
+};
+initObj["objField"]["newSubPath"] = {
+    b: valueWithNestedObject
+};
+assertAddFieldsResult({"objField.newSubPath.b": valueWithNestedObject}, initObj);
+assertAddFieldsResult({"objField.newSubPath.b": {$literal: valueWithNestedObject}}, initObj);
+initObj["objField"] = {
+    p: {q: 1},
+    subArr: [1]
+};  // Reset input object.
+
+// When the top level field doesn"t exist, a new nested object is created based on the given path.
+initObj["newField"] = {
+    newSubPath: {b: "newValue"}
+};
+assertAddFieldsResult({"newField.newSubPath.b": {$literal: "newValue"}}, initObj);
+assertAddFieldsResult({"newField.newSubPath.b": "newValue"}, initObj);
+
+// When the top level field doesn"t exist, a new nested object is created based on the given path
+// and the structure of the object in the value.
+initObj["newField"]["newSubPath"] = {
+    b: valueWithNestedObject
+};
+assertAddFieldsResult({"newField.newSubPath.b": valueWithNestedObject}, initObj);
+assertAddFieldsResult({"newField.newSubPath.b": {$literal: valueWithNestedObject}}, initObj);
+delete initObj["newField"];  // Reset.
+
+// Test when the path encounters an array and the value is a scalar.
+initObj["arrayField"] = {
+    newSubPath: {b: "newValue"}
+};
+let expectedSubObj = {newSubPath: {b: "newValue"}};
+initObj["arrayField"] =
+    [expectedSubObj, Object.assign({subField: [2, {}]}, expectedSubObj), [expectedSubObj]];
+assertAddFieldsResult({"arrayField.newSubPath.b": {$literal: "newValue"}}, initObj);
+assertAddFieldsResult({"arrayField.newSubPath.b": "newValue"}, initObj);
+
+// Test when the path encounters an array and the value is a nested object.
+expectedSubObj = {
+    newSubPath: {b: valueWithNestedObject}
+};
+initObj["arrayField"] =
+    [expectedSubObj, Object.assign({subField: [2, {}]}, expectedSubObj), [expectedSubObj]];
+assertAddFieldsResult({"arrayField.newSubPath.b": valueWithNestedObject}, initObj);
+assertAddFieldsResult({"arrayField.newSubPath.b": {$literal: valueWithNestedObject}}, initObj);
+
+// Test when the path encounters multiple arrays and the value is a nested object.
+expectedSubObj = {
+    subField: {b: valueWithNestedObject}
+};
+initObj["arrayField"] = [
+    expectedSubObj,
+    {
+        subField:
+            [{b: valueWithNestedObject}, {b: valueWithNestedObject}]  // Sub-array is also exploded.
+    },
+    [expectedSubObj]
+];
+assertAddFieldsResult({"arrayField.subField.b": valueWithNestedObject}, initObj);
+assertAddFieldsResult({"arrayField.subField.b": {$literal: valueWithNestedObject}}, initObj);
+})();
diff --git a/src/mongo/db/commands/mr_common.cpp b/src/mongo/db/commands/mr_common.cpp
index b0698082778..64b8337bc64 100644
--- a/src/mongo/db/commands/mr_common.cpp
+++ b/src/mongo/db/commands/mr_common.cpp
@@ -111,7 +111,7 @@ auto translateMap(boost::intrusive_ptr<ExpressionContext> expCtx, std::string co
     auto emitExpression = ExpressionInternalJsEmit::create(
         expCtx, ExpressionFieldPath::parse(expCtx, "$$ROOT", expCtx->variablesParseState), code);
     auto node = std::make_unique<projection_executor::InclusionNode>(
-        ProjectionPolicies{ProjectionPolicies::DefaultIdPolicy::kExcludeId});
+        ProjectionPolicies{ProjectionPolicies::DefaultIdPolicy::kExcludeId}, false);
     node->addExpressionForPath(FieldPath{"emits"s}, std::move(emitExpression));
     auto inclusion = std::unique_ptr<TransformerInterface>{
         std::make_unique<projection_executor::InclusionProjectionExecutor>(
@@ -153,7 +153,7 @@ auto translateFinalize(boost::intrusive_ptr<ExpressionContext> expCtx,
             code,
             ExpressionFunction::kJavaScript);
         auto node = std::make_unique<projection_executor::InclusionNode>(
-            ProjectionPolicies{ProjectionPolicies::DefaultIdPolicy::kIncludeId});
+            ProjectionPolicies{ProjectionPolicies::DefaultIdPolicy::kIncludeId}, false);
         node->addProjectionForPath(FieldPath{"_id"s});
         node->addExpressionForPath(FieldPath{"value"s}, std::move(jsExpression));
         auto inclusion = std::unique_ptr<TransformerInterface>{
diff --git a/src/mongo/db/exec/add_fields_projection_executor.cpp b/src/mongo/db/exec/add_fields_projection_executor.cpp
index ec3acc136bc..962666132bc 100644
--- a/src/mongo/db/exec/add_fields_projection_executor.cpp
+++ b/src/mongo/db/exec/add_fields_projection_executor.cpp
@@ -280,14 +280,14 @@ void AddFieldsProjectionExecutor::parseSubObject(const BSONObj& subObj,
                                                  const VariablesParseState& variablesParseState,
                                                  InclusionNode* node) {
     for (auto&& elem : subObj) {
-        invariant(elem.fieldName()[0] != '$');
+        std::string fieldName = elem.fieldName();
+        invariant(fieldName[0] != '$');
         // Dotted paths in a sub-object have already been detected and disallowed by the function
         // ProjectionSpecValidator::validate().
-        invariant(elem.fieldNameStringData().find('.') == std::string::npos);
+        invariant(fieldName.find('.') == std::string::npos);
 
         if (elem.type() == BSONType::Object) {
             // This is either an expression, or a nested specification.
-            auto fieldName = elem.fieldNameStringData().toString();
             if (!parseObjectAsExpression(
                     FieldPath::getFullyQualifiedPath(node->getPath(), fieldName),
                     elem.Obj(),
diff --git a/src/mongo/db/exec/add_fields_projection_executor.h b/src/mongo/db/exec/add_fields_projection_executor.h
index 00ba580c2c7..eb1f414280c 100644
--- a/src/mongo/db/exec/add_fields_projection_executor.h
+++ b/src/mongo/db/exec/add_fields_projection_executor.h
@@ -55,7 +55,7 @@ public:
                              {ProjectionPolicies::DefaultIdPolicy::kIncludeId,
                               ProjectionPolicies::ArrayRecursionPolicy::kRecurseNestedArrays,
                               ProjectionPolicies::ComputedFieldsPolicy::kAllowComputedFields}),
-          _root(new InclusionNode(_policies)) {}
+          _root(new InclusionNode(_policies, true)) {}
 
     /**
      * Creates the data needed to perform an AddFields.
diff --git a/src/mongo/db/exec/exclusion_projection_executor.h b/src/mongo/db/exec/exclusion_projection_executor.h
index cf21bf57126..b7767175c66 100644
--- a/src/mongo/db/exec/exclusion_projection_executor.h
+++ b/src/mongo/db/exec/exclusion_projection_executor.h
@@ -44,7 +44,7 @@ namespace mongo::projection_executor {
 class ExclusionNode final : public ProjectionNode {
 public:
     ExclusionNode(ProjectionPolicies policies, std::string pathToNode = "")
-        : ProjectionNode(policies, std::move(pathToNode)) {}
+        : ProjectionNode(policies, false, std::move(pathToNode)) {}
 
     ExclusionNode* addOrGetChild(const std::string& field) {
         return static_cast<ExclusionNode*>(ProjectionNode::addOrGetChild(field));
diff --git a/src/mongo/db/exec/inclusion_projection_executor.h b/src/mongo/db/exec/inclusion_projection_executor.h
index b095a423652..b26b62859ed 100644
--- a/src/mongo/db/exec/inclusion_projection_executor.h
+++ b/src/mongo/db/exec/inclusion_projection_executor.h
@@ -43,8 +43,10 @@ namespace mongo::projection_executor {
  */
 class InclusionNode : public ProjectionNode {
 public:
-    InclusionNode(ProjectionPolicies policies, std::string pathToNode = "")
-        : ProjectionNode(policies, std::move(pathToNode)) {}
+    InclusionNode(ProjectionPolicies policies,
+                  bool containsComputedFields,
+                  std::string pathToNode = "")
+        : ProjectionNode(policies, containsComputedFields, std::move(pathToNode)) {}
 
     InclusionNode* addOrGetChild(const std::string& field) {
         return static_cast<InclusionNode*>(ProjectionNode::addOrGetChild(field));
@@ -83,7 +85,9 @@ protected:
     }
     std::unique_ptr<ProjectionNode> makeChild(const std::string& fieldName) const override {
         return std::make_unique<InclusionNode>(
-            _policies, FieldPath::getFullyQualifiedPath(_pathToNode, fieldName));
+            _policies,
+            _subtreeContainsComputedFields,
+            FieldPath::getFullyQualifiedPath(_pathToNode, fieldName));
     }
     MutableDocument initializeOutputDocument(const Document& inputDoc) const final {
         // Technically this value could be min(number of projected fields, size of input
@@ -112,15 +116,19 @@ protected:
  */
 class FastPathEligibleInclusionNode final : public InclusionNode {
 public:
-    FastPathEligibleInclusionNode(ProjectionPolicies policies, std::string pathToNode = "")
-        : InclusionNode(policies, std::move(pathToNode)) {}
+    FastPathEligibleInclusionNode(ProjectionPolicies policies,
+                                  bool containsComputedFields,
+                                  std::string pathToNode = "")
+        : InclusionNode(policies, containsComputedFields, std::move(pathToNode)) {}
 
     Document applyToDocument(const Document& inputDoc) const final;
 
 protected:
     std::unique_ptr<ProjectionNode> makeChild(const std::string& fieldName) const final {
         return std::make_unique<FastPathEligibleInclusionNode>(
-            _policies, FieldPath::getFullyQualifiedPath(_pathToNode, fieldName));
+            _policies,
+            _subtreeContainsComputedFields,
+            FieldPath::getFullyQualifiedPath(_pathToNode, fieldName));
     }
 
 private:
@@ -147,8 +155,8 @@ public:
         : InclusionProjectionExecutor(
               expCtx,
               policies,
-              allowFastPath ? std::make_unique<FastPathEligibleInclusionNode>(policies)
-                            : std::make_unique<InclusionNode>(policies)) {}
+              allowFastPath ? std::make_unique<FastPathEligibleInclusionNode>(policies, false)
+                            : std::make_unique<InclusionNode>(policies, false)) {}
 
     TransformerType getType() const final {
         return TransformerType::kInclusionProjection;
diff --git a/src/mongo/db/exec/projection_executor_test.cpp b/src/mongo/db/exec/projection_executor_test.cpp
index 3c6fea2808e..51642430faf 100644
--- a/src/mongo/db/exec/projection_executor_test.cpp
+++ b/src/mongo/db/exec/projection_executor_test.cpp
@@ -573,7 +573,7 @@ TEST(ProjectionExecutorType, GetExpressionForPathGetsTopLevelExpression) {
     auto projectObj = BSON("$add" << BSON_ARRAY(BSON("$const" << 1) << BSON("$const" << 3)));
     auto expr = Expression::parseObject(expCtx, projectObj, expCtx->variablesParseState);
     ProjectionPolicies defaultPolicies;
-    auto node = InclusionNode(defaultPolicies);
+    auto node = InclusionNode(defaultPolicies, false);
     node.addExpressionForPath(FieldPath("key"), expr);
     BSONObjBuilder bob;
     ASSERT_EQ(expr, node.getExpressionForPath(FieldPath("key")));
@@ -586,7 +586,7 @@ TEST(ProjectionExecutorType, GetExpressionForPathGetsCorrectTopLevelExpression)
     auto correctExpr = Expression::parseObject(expCtx, correctObj, expCtx->variablesParseState);
     auto incorrectExpr = Expression::parseObject(expCtx, incorrectObj, expCtx->variablesParseState);
     ProjectionPolicies defaultPolicies;
-    auto node = InclusionNode(defaultPolicies);
+    auto node = InclusionNode(defaultPolicies, false);
     node.addExpressionForPath(FieldPath("key"), correctExpr);
     node.addExpressionForPath(FieldPath("other"), incorrectExpr);
     BSONObjBuilder bob;
@@ -598,7 +598,7 @@ TEST(ProjectionExecutorType, GetExpressionForPathGetsNonTopLevelExpression) {
     auto projectObj = BSON("$add" << BSON_ARRAY(BSON("$const" << 1) << BSON("$const" << 3)));
     auto expr = Expression::parseObject(expCtx, projectObj, expCtx->variablesParseState);
     ProjectionPolicies defaultPolicies;
-    auto node = InclusionNode(defaultPolicies);
+    auto node = InclusionNode(defaultPolicies, false);
     node.addExpressionForPath(FieldPath("key.second"), expr);
     BSONObjBuilder bob;
     ASSERT_EQ(expr, node.getExpressionForPath(FieldPath("key.second")));
diff --git a/src/mongo/db/exec/projection_node.cpp b/src/mongo/db/exec/projection_node.cpp
index 9ea6e9afbf9..8763c087810 100644
--- a/src/mongo/db/exec/projection_node.cpp
+++ b/src/mongo/db/exec/projection_node.cpp
@@ -36,8 +36,12 @@ using ArrayRecursionPolicy = ProjectionPolicies::ArrayRecursionPolicy;
 using ComputedFieldsPolicy = ProjectionPolicies::ComputedFieldsPolicy;
 using DefaultIdPolicy = ProjectionPolicies::DefaultIdPolicy;
 
-ProjectionNode::ProjectionNode(ProjectionPolicies policies, std::string pathToNode)
-    : _policies(policies), _pathToNode(std::move(pathToNode)) {}
+ProjectionNode::ProjectionNode(ProjectionPolicies policies,
+                               bool containsComputedFields,
+                               std::string pathToNode)
+    : _policies(policies),
+      _pathToNode(std::move(pathToNode)),
+      _subtreeContainsComputedFields(containsComputedFields) {}
 
 void ProjectionNode::addProjectionForPath(const FieldPath& path) {
     makeOptimizationsStale();
diff --git a/src/mongo/db/exec/projection_node.h b/src/mongo/db/exec/projection_node.h
index 09c191226a1..a12d3c04ca2 100644
--- a/src/mongo/db/exec/projection_node.h
+++ b/src/mongo/db/exec/projection_node.h
@@ -47,7 +47,9 @@ namespace mongo::projection_executor {
  */
 class ProjectionNode {
 public:
-    ProjectionNode(ProjectionPolicies policies, std::string pathToNode = "");
+    ProjectionNode(ProjectionPolicies policies,
+                   bool containsComputedFields,
+                   std::string pathToNode = "");
 
     virtual ~ProjectionNode() = default;
 
@@ -160,7 +162,7 @@ protected:
     std::string _pathToNode;
 
     // Whether this node or any child of this node contains a computed field.
-    bool _subtreeContainsComputedFields{false};
+    bool _subtreeContainsComputedFields;
 
 private:
     // Iterates 'inputDoc' for each projected field, adding to or removing from 'outputDoc'. Also