SERVER-67250 Do not overwrite existing objects in 'addToArray()'

4 files changed, 93 insertions, 15 deletions

diff --git a/jstests/core/columnstore_index_correctness.js b/jstests/core/columnstore_index_correctness.js
index b0115af1398..324b32def1e 100644
--- a/jstests/core/columnstore_index_correctness.js
+++ b/jstests/core/columnstore_index_correctness.js
@@ -511,6 +511,7 @@ const docs = [
     {a: {b: [{c: [1, 2]}]}},
     {a: {b: {c: [1, 2]}}},
     {a: [[1, 2], [{b: [[1, 2], [{c: [[1, 2], [{}], 2]}], 2]}], 2]},
+    {a: [{m: 1, n: 2}, {m: 2, o: 1}]},
 ];
 
 let docNum = 0;
@@ -537,7 +538,7 @@ const kProjection = {
 };
 
 // Run an explain.
-const explain = coll.find({}, kProjection).explain();
+let explain = coll.find({}, kProjection).explain();
 assert(planHasStage(db, explain, "COLUMN_SCAN"), explain);
 
 // Run a query getting all of the results using the column index.
@@ -549,4 +550,24 @@ for (let res of results) {
     const originalDoc = coll.findOne({num: res.num});
     assert.eq(res, trueResult, originalDoc);
 }
+
+// Run a similar query that projects multiple fields with a shared parent object.
+const kSiblingProjection = {
+    _id: 0,
+    "a.m": 1,
+    "a.n": 1,
+    num: 1
+};
+
+explain = coll.find({}, kSiblingProjection).explain();
+assert(planHasStage(db, explain, "COLUMN_SCAN"), explain);
+
+results = coll.find({}, kSiblingProjection).toArray();
+assert.gt(results.length, 0);
+for (let res of results) {
+    const trueResult =
+        coll.find({num: res.num}, kSiblingProjection).hint({$natural: 1}).toArray()[0];
+    const originalDoc = coll.findOne({num: res.num});
+    assert.eq(res, trueResult, originalDoc);
+}
 })();
diff --git a/jstests/core/projection_semantics.js b/jstests/core/projection_semantics.js
index a9b1c32bf01..42605b3df49 100644
--- a/jstests/core/projection_semantics.js
+++ b/jstests/core/projection_semantics.js
@@ -201,6 +201,42 @@ function testInputOutput({input, projection, expectedOutput, interestingIndexes
             {b: 1},
         ]
     });
+
+    // Test the case where two paths in a projection go through the same parent object.
+    const testIncludeOnlyADotBAndADotC = (input, output) => testInputOutput({
+        input: input,
+        projection: {'a.b': 1, 'a.c': 1, _id: 0},
+        expectedOutput: output,
+        interestingIndexes:
+            [{a: 1}, {'a.b': 1}, {'a.c': 1}, {'a.b': 1, 'a.c': 1}, {'a.b': 1, 'a.c': -1}]
+    });
+    testIncludeOnlyADotBAndADotC({_id: 0, a: {b: "scalar", c: "scalar", d: "extra"}},
+                                 {a: {b: "scalar", c: "scalar"}});
+    testIncludeOnlyADotBAndADotC({_id: 1, a: [{b: 1, c: 2, d: 3}, {b: 4, c: 5, d: 6}]},
+                                 {a: [{b: 1, c: 2}, {b: 4, c: 5}]});
+
+    // Array cases where one or both of the paths don't exist.
+    testIncludeOnlyADotBAndADotC({_id: 5, a: [{b: 1, c: 2}, {b: 3, d: 4}]},
+                                 {a: [{b: 1, c: 2}, {b: 3}]});
+    testIncludeOnlyADotBAndADotC({_id: 6, a: [{c: 1, d: 2}, {b: 3, d: 4}]}, {a: [{c: 1}, {b: 3}]});
+    testIncludeOnlyADotBAndADotC({_id: 7, a: []}, {a: []});
+    testIncludeOnlyADotBAndADotC({_id: 8, a: [{b: 1, c: 2}, "extra", {b: 3, c: 4}]},
+                                 {a: [{b: 1, c: 2}, {b: 3, c: 4}]});
+
+    // Non-array cases where one or both of the paths don't exist.
+    //
+    // TODO SERVER-23229: This will return different results if there is a covering index, so here
+    // but not elsewhere we don't use any "interestingIndexes" in test cases.
+    const testIncludeADotBAndCNoIndexes = (input, output) => testInputOutput({
+        input: input,
+        projection: {'a.b': 1, 'a.c': 1, _id: 0},
+        expectedOutput: output,
+        interestingIndexes: []
+    });
+
+    testIncludeADotBAndCNoIndexes({_id: 2, a: {b: "scalar", d: "extra"}}, {a: {b: "scalar"}});
+    testIncludeADotBAndCNoIndexes({_id: 3, a: {c: "scalar", d: "extra"}}, {a: {c: "scalar"}});
+    testIncludeADotBAndCNoIndexes({_id: 4, a: {d: "extra"}}, {a: {}});
 }());
 
 (function testInclusionLevelsOfNesting() {
diff --git a/src/mongo/db/exec/sbe/values/columnar.cpp b/src/mongo/db/exec/sbe/values/columnar.cpp
index 7490d549803..c1bd51f6b69 100644
--- a/src/mongo/db/exec/sbe/values/columnar.cpp
+++ b/src/mongo/db/exec/sbe/values/columnar.cpp
@@ -237,6 +237,24 @@ void addToObjectNoArrays(value::TypeTags tag,
     });
 }
 
+/*
+ * Ensures that the path (stored in 'state') leads to an object and materializes an empty object if
+ * it does not. Assumes that there are no arrays along remaining path (i.e., the components that are
+ * not yet traversed via withNextPathComponent()).
+ *
+ * This function is a no-op when there are no remaining path components.
+ */
+template <class C>
+void materializeObjectNoArrays(AddToDocumentState<C>& state, value::Object& out) {
+    if (state.atLastPathComponent()) {
+        return;
+    }
+
+    state.withNextPathComponent([&](StringData nextPathComponent) {
+        materializeObjectNoArrays(state, *findOrAddObjInObj(nextPathComponent, &out));
+    });
+}
+
 template <class C>
 void addToObject(value::Object& obj, AddToDocumentState<C>& state);
 
@@ -268,23 +286,19 @@ void addToArray(value::Array& arr, AddToDocumentState<C>& state) {
                 for (; insertAt < index; insertAt++) {
                     invariant(insertAt < arr.size());
 
-                    auto [tag, val] = [nextChar, &state]() {
-                        if (nextChar == '|') {
-                            return state.extractAndCopyValue();
+                    if (nextChar == 'o') {
+                        materializeObjectNoArrays(state, *findOrAddObjInArr(insertAt, &arr));
+                    } else if (nextChar == '|') {
+                        auto [tag, val] = state.extractAndCopyValue();
+                        if (state.atLastPathComponent()) {
+                            invariant(arr.getAt(insertAt).first == kPlaceHolderType);
+                            arr.setAt(insertAt, tag, val);
                         } else {
-                            invariant(nextChar == 'o');
-                            return value::makeNewObject();
+                            addToObjectNoArrays(
+                                tag, val, state, *findOrAddObjInArr(insertAt, &arr), 0);
                         }
-                    }();
-                    if (state.atLastPathComponent()) {
-                        // At this point we are inserting a leaf value.
-                        dassert(arr.getAt(insertAt).first == kPlaceHolderType);
-                        arr.setAt(insertAt, tag, val);
                     } else {
-                        // This is valid on initialized elements when the subobject contains more
-                        // than one member.
-                        auto* subObj = findOrAddObjInArr(insertAt, &arr);
-                        addToObjectNoArrays(tag, val, state, *subObj, 0);
+                        MONGO_UNREACHABLE;
                     }
                 }
                 break;
diff --git a/src/mongo/db/exec/sbe/values/columnar_test.cpp b/src/mongo/db/exec/sbe/values/columnar_test.cpp
index 9dc9e7717d0..ebbed88848a 100644
--- a/src/mongo/db/exec/sbe/values/columnar_test.cpp
+++ b/src/mongo/db/exec/sbe/values/columnar_test.cpp
@@ -201,4 +201,11 @@ TEST(ColumnarObjTest, AddNonLeafCellWithArrayInfoToObject) {
     std::vector<MockTranslatedCell> cells{makeCellOfIntegers("a.b", "{[o1", {})};
     compareMakeObjWithExpected(cells, fromjson("{a: {b: [{}, {}]}}"));
 }
+
+TEST(ColumnarObjTest, AddLeafCellThenAddSparseSibling) {
+    std::vector<MockTranslatedCell> cells{makeCellOfIntegers("a.b", "[", {1, 2}),
+                                          makeCellOfIntegers("a", "[o1", {}),
+                                          makeCellOfIntegers("a.c", "[1", {3})};
+    compareMakeObjWithExpected(cells, fromjson("{a: [{b: 1}, {b: 2, c: 3}]}"));
+}
 }  // namespace mongo::sbe