SERVER-61997 Incompatible object with only empty subobj should not re-start subobj compression

3 files changed, 91 insertions, 6 deletions

diff --git a/src/mongo/bson/util/bsoncolumn_test.cpp b/src/mongo/bson/util/bsoncolumn_test.cpp
index 2ddcca2425b..e3656b7f059 100644
--- a/src/mongo/bson/util/bsoncolumn_test.cpp
+++ b/src/mongo/bson/util/bsoncolumn_test.cpp
@@ -3679,6 +3679,74 @@ TEST_F(BSONColumnTest, ObjectWithOnlyEmptyObjsDoesNotStartInterleaving) {
     verifyDecompression(binData, elems);
 }
 
+TEST_F(BSONColumnTest, ObjectWithOnlyEmptyObjsDoesNotStartInterleavingFromDetermine) {
+    BSONColumnBuilder cb("test"_sd);
+
+    // Append elements so we are in kSubObjDeterminingReference state when element with 'b' field is
+    // appended. Make sure this does not re-start subobj compression as it only contain empty
+    // subobj.
+    std::vector<BSONElement> elems;
+    elems.push_back(createElementObj(BSON("a" << 1)));
+    elems.push_back(createElementObj(BSON("b" << BSONObjBuilder().obj())));
+
+    for (const auto& elem : elems) {
+        cb.append(elem);
+    }
+
+    BufBuilder expected;
+    appendInterleavedStart(expected, elems.front().Obj());
+    appendSimple8bControl(expected, 0b1000, 0b0000);
+    appendSimple8bBlock64(expected, kDeltaForBinaryEqualValues);
+    appendEOO(expected);
+
+    appendLiteral(expected, elems[1]);
+    appendEOO(expected);
+
+    auto binData = cb.finalize();
+    verifyBinary(binData, expected);
+    verifyDecompression(binData, elems);
+}
+
+TEST_F(BSONColumnTest, ObjectWithOnlyEmptyObjsDoesNotStartInterleavingFromAppending) {
+    BSONColumnBuilder cb("test"_sd);
+
+    // Append enough elements so we are in kSubObjAppending state when element with 'b' field is
+    // appended. Make sure this does not re-start subobj compression as it only contain empty
+    // subobj.
+    std::vector<BSONElement> elems;
+    elems.push_back(createElementObj(BSON("a" << 1)));
+    elems.push_back(createElementObj(BSON("a" << 2)));
+    elems.push_back(createElementObj(BSON("a" << 3)));
+    elems.push_back(createElementObj(BSON("a" << 4)));
+    elems.push_back(createElementObj(BSON("a" << 5)));
+    elems.push_back(createElementObj(BSON("a" << 6)));
+    elems.push_back(createElementObj(BSON("b" << BSONObjBuilder().obj())));
+
+    for (const auto& elem : elems) {
+        cb.append(elem);
+    }
+
+    BufBuilder expected;
+    appendInterleavedStart(expected, elems.front().Obj());
+    appendSimple8bControl(expected, 0b1000, 0b0000);
+    appendSimple8bBlocks64(expected,
+                           {kDeltaForBinaryEqualValues,
+                            deltaInt32(elems[1].Obj()["a"], elems[0].Obj()["a"]),
+                            deltaInt32(elems[2].Obj()["a"], elems[1].Obj()["a"]),
+                            deltaInt32(elems[3].Obj()["a"], elems[2].Obj()["a"]),
+                            deltaInt32(elems[4].Obj()["a"], elems[3].Obj()["a"]),
+                            deltaInt32(elems[5].Obj()["a"], elems[4].Obj()["a"])},
+                           1);
+    appendEOO(expected);
+
+    appendLiteral(expected, elems[6]);
+    appendEOO(expected);
+
+    auto binData = cb.finalize();
+    verifyBinary(binData, expected);
+    verifyDecompression(binData, elems);
+}
+
 TEST_F(BSONColumnTest, NonZeroRLEInFirstBlockAfterSimple8bBlocks) {
     BSONColumnBuilder cb("test"_sd);
 
diff --git a/src/mongo/bson/util/bsoncolumnbuilder.cpp b/src/mongo/bson/util/bsoncolumnbuilder.cpp
index 88d60603f9f..139a7adf747 100644
--- a/src/mongo/bson/util/bsoncolumnbuilder.cpp
+++ b/src/mongo/bson/util/bsoncolumnbuilder.cpp
@@ -342,6 +342,13 @@ BSONColumnBuilder& BSONColumnBuilder::append(BSONElement elem) {
                 // If merge failed, flush current sub-object compression and start over.
                 _flushSubObjMode();
 
+                // If we only contain empty subobj (no value elements) then append in regular mode
+                // instead of re-starting subobj compression.
+                if (numElements == 0) {
+                    _state.append(elem);
+                    return *this;
+                }
+
                 _referenceSubObj = obj.getOwned();
                 _bufferedObjElements.push_back(_referenceSubObj);
                 _mode = Mode::kSubObjDeterminingReference;
@@ -361,7 +368,15 @@ BSONColumnBuilder& BSONColumnBuilder::append(BSONElement elem) {
     }
 
     // Reference already determined for sub-object compression, try to add this new object.
-    _appendSubElements(obj);
+    if (!_appendSubElements(obj)) {
+        // If we were not compatible restart subobj compression unless our object contain no value
+        // fields (just empty subobjects)
+        if (numElements == 0) {
+            _state.append(elem);
+        } else {
+            _startDetermineSubObjReference(obj);
+        }
+    }
     return *this;
 }
 
@@ -903,7 +918,7 @@ Simple8bWriteFn BSONColumnBuilder::EncodingState::_createBufferWriter() {
     };
 }
 
-void BSONColumnBuilder::_appendSubElements(const BSONObj& obj) {
+bool BSONColumnBuilder::_appendSubElements(const BSONObj& obj) {
     // Check if added object is compatible with selected reference object. Collect a flat vector of
     // all elements while we are doing this.
     _flattenedAppendedObj.clear();
@@ -912,8 +927,7 @@ void BSONColumnBuilder::_appendSubElements(const BSONObj& obj) {
                 _flattenedAppendedObj.push_back(elem);
             })) {
         _flushSubObjMode();
-        _startDetermineSubObjReference(obj);
-        return;
+        return false;
     }
 
     // We should have recieved one callback for every sub-element in reference object. This should
@@ -932,6 +946,7 @@ void BSONColumnBuilder::_appendSubElements(const BSONObj& obj) {
         else
             state.skip();
     }
+    return true;
 }
 
 void BSONColumnBuilder::_startDetermineSubObjReference(const BSONObj& obj) {
@@ -989,7 +1004,9 @@ void BSONColumnBuilder::_finishDetermineSubObjReference() {
     auto it = _bufferedObjElements.begin() + 1;
     auto end = _bufferedObjElements.end();
     for (; it != end; ++it) {
-        _appendSubElements(*it);
+        // The objects we append here should always be compatible with our reference object. If they
+        // are not then there is a bug somewhere.
+        invariant(_appendSubElements(*it));
     }
     _bufferedObjElements.clear();
 }
diff --git a/src/mongo/bson/util/bsoncolumnbuilder.h b/src/mongo/bson/util/bsoncolumnbuilder.h
index 28cea020ca5..81e4e8c6416 100644
--- a/src/mongo/bson/util/bsoncolumnbuilder.h
+++ b/src/mongo/bson/util/bsoncolumnbuilder.h
@@ -152,7 +152,7 @@ private:
     };
 
     // Append Object for sub-object compression when in mode kSubObjAppending
-    void _appendSubElements(const BSONObj& obj);
+    bool _appendSubElements(const BSONObj& obj);
 
     // Transition into kSubObjDeterminingReference mode
     void _startDetermineSubObjReference(const BSONObj& obj);