SERVER-68094: Use $replaceRoot instead of $project

(cherry picked from commit 015dc2badcafc3238845b0eec3d6084fdff2545c)

3 files changed, 86 insertions, 45 deletions

diff --git a/src/mongo/db/pipeline/resharding_initial_split_policy_test.cpp b/src/mongo/db/pipeline/resharding_initial_split_policy_test.cpp
index a008cc92115..ccaecddf502 100644
--- a/src/mongo/db/pipeline/resharding_initial_split_policy_test.cpp
+++ b/src/mongo/db/pipeline/resharding_initial_split_policy_test.cpp
@@ -35,6 +35,7 @@
 #include "mongo/db/pipeline/document_source_mock.h"
 #include "mongo/db/pipeline/sharded_agg_helpers.h"
 #include "mongo/db/s/config/initial_split_policy.h"
+#include "mongo/logv2/log.h"
 #include "mongo/s/query/sharded_agg_test_fixture.h"
 #include "mongo/unittest/unittest.h"
 
@@ -47,7 +48,6 @@ const ShardId primaryShardId = ShardId("0");
 
 TEST_F(ReshardingSplitPolicyTest, ShardKeyWithNonDottedFieldAndIdIsNotProjectedSucceeds) {
     auto shardKeyPattern = ShardKeyPattern(BSON("a" << 1));
-
     auto pipeline =
         Pipeline::parse(ReshardingSplitPolicy::createRawPipeline(
                             shardKeyPattern, 2 /* samplingRatio */, 1 /* numSplitPoints */),
@@ -55,7 +55,6 @@ TEST_F(ReshardingSplitPolicyTest, ShardKeyWithNonDottedFieldAndIdIsNotProjectedS
     auto mockSource =
         DocumentSourceMock::createForTest({"{_id: 10, a: 15}", "{_id: 3, a: 5}"}, expCtx());
     pipeline->addInitialSource(mockSource.get());
-
     // We sample all of the documents since numSplitPoints(1) * samplingRatio (2) = 2 and the
     // document source has 2 chunks. So we can assert on the returned values.
     auto next = pipeline->getNext();
@@ -92,7 +91,6 @@ TEST_F(ReshardingSplitPolicyTest, ShardKeyWithIdFieldIsProjectedSucceeds) {
 TEST_F(ReshardingSplitPolicyTest, CompoundShardKeyWithNonDottedHashedFieldSucceeds) {
     auto shardKeyPattern = ShardKeyPattern(BSON("a" << 1 << "b"
                                                     << "hashed"));
-
     auto pipeline =
         Pipeline::parse(ReshardingSplitPolicy::createRawPipeline(
                             shardKeyPattern, 2 /* samplingRatio */, 1 /* numSplitPoints */),
@@ -100,7 +98,6 @@ TEST_F(ReshardingSplitPolicyTest, CompoundShardKeyWithNonDottedHashedFieldSuccee
     auto mockSource = DocumentSourceMock::createForTest(
         {"{x: 1, b: 16, a: 15}", "{x: 2, b: 123, a: 5}"}, expCtx());
     pipeline->addInitialSource(mockSource.get());
-
     // We sample all of the documents since numSplitPoints(1) * samplingRatio (2) = 2 and the
     // document source has 2 chunks. So we can assert on the returned values.
     auto next = pipeline->getNext();
@@ -128,9 +125,9 @@ TEST_F(ReshardingSplitPolicyTest, CompoundShardKeyWithDottedFieldSucceeds) {
     // We sample all of the documents since numSplitPoints(1) * samplingRatio (2) = 2 and the
     // document source has 2 chunks. So we can assert on the returned values.
     auto next = pipeline->getNext();
-    ASSERT_BSONOBJ_EQ(next.get().toBson(), BSON("a" << BSON("b" << 10) << "c" << 5));
+    ASSERT_BSONOBJ_EQ(next.value().toBson(), BSON("a.b" << 10 << "c" << 5));
     next = pipeline->getNext();
-    ASSERT_BSONOBJ_EQ(next.get().toBson(), BSON("a" << BSON("b" << 20) << "c" << 1));
+    ASSERT_BSONOBJ_EQ(next.value().toBson(), BSON("a.b" << 20 << "c" << 1));
     ASSERT(!pipeline->getNext());
 }
 
@@ -149,11 +146,11 @@ TEST_F(ReshardingSplitPolicyTest, CompoundShardKeyWithDottedHashedFieldSucceeds)
     // We sample all of the documents since numSplitPoints(1) * samplingRatio (2) = 2 and the
     // document source has 2 chunks. So we can assert on the returned values.
     auto next = pipeline->getNext();
-    ASSERT_BSONOBJ_EQ(next.get().toBson(),
-                      BSON("a" << BSON("b" << 10 << "c" << -6548868637522515075LL) << "c" << 5));
+    ASSERT_BSONOBJ_EQ(next.value().toBson(),
+                      BSON("a.b" << 10 << "c" << 5 << "a.c" << -6548868637522515075LL));
     next = pipeline->getNext();
-    ASSERT_BSONOBJ_EQ(next.get().toBson(),
-                      BSON("a" << BSON("b" << 20 << "c" << 2598032665634823220LL) << "c" << 1));
+    ASSERT_BSONOBJ_EQ(next.value().toBson(),
+                      BSON("a.b" << 20 << "c" << 1 << "a.c" << 2598032665634823220LL));
     ASSERT(!pipeline->getNext());
 }
 
@@ -200,5 +197,57 @@ TEST_F(ReshardingSplitPolicyTest, SamplingSuceeds) {
     }
 }
 
+TEST_F(ReshardingSplitPolicyTest, ShardKeyWithDottedPathAndIdIsNotProjectedSucceeds) {
+    auto shardKeyPattern = ShardKeyPattern(BSON("b" << 1));
+    auto pipeline =
+        Pipeline::parse(ReshardingSplitPolicy::createRawPipeline(
+                            shardKeyPattern, 2 /* samplingRatio */, 1 /* numSplitPoints */),
+                        expCtx());
+    auto mockSource = DocumentSourceMock::createForTest(
+        {"{_id: {a: 15}, b: 10}", "{_id: {a: 5}, b:1}"}, expCtx());
+    pipeline->addInitialSource(mockSource.get());
+    auto next = pipeline->getNext();
+    ASSERT_BSONOBJ_EQ(next.value().toBson(), BSON("b" << 1));
+    next = pipeline->getNext();
+    ASSERT_BSONOBJ_EQ(next.value().toBson(), BSON("b" << 10));
+    ASSERT(!pipeline->getNext());
+}
+
+TEST_F(ReshardingSplitPolicyTest, CompoundShardKeyWithDottedPathAndIdIsProjectedSucceeds) {
+    auto shardKeyPattern = ShardKeyPattern(BSON("_id.a" << 1 << "c" << 1));
+    auto pipeline =
+        Pipeline::parse(ReshardingSplitPolicy::createRawPipeline(
+                            shardKeyPattern, 2 /* samplingRatio */, 1 /* numSplitPoints */),
+                        expCtx());
+    auto mockSource = DocumentSourceMock::createForTest(
+        {"{_id: {a: 15}, c: 10}", "{_id: {a: 5}, c: 1}"}, expCtx());
+    pipeline->addInitialSource(mockSource.get());
+    auto next = pipeline->getNext();
+    ASSERT_BSONOBJ_EQ(next.value().toBson(), BSON("_id.a" << 5 << "c" << 1));
+    next = pipeline->getNext();
+    ASSERT_BSONOBJ_EQ(next.value().toBson(), BSON("_id.a" << 15 << "c" << 10));
+    ASSERT(!pipeline->getNext());
+}
+
+TEST_F(ReshardingSplitPolicyTest, CompoundShardKeyWithDottedHashedPathSucceeds) {
+    auto shardKeyPattern = ShardKeyPattern(BSON("_id.a" << 1 << "b" << 1 << "_id.b"
+                                                        << "hashed"));
+
+    auto pipeline =
+        Pipeline::parse(ReshardingSplitPolicy::createRawPipeline(
+                            shardKeyPattern, 2 /* samplingRatio */, 1 /* numSplitPoints */),
+                        expCtx());
+    auto mockSource = DocumentSourceMock::createForTest(
+        {"{x: 10, _id: {a: 20, b: 16}, b: 1}", "{x: 3, _id: {a: 10, b: 123}, b: 5}"}, expCtx());
+    pipeline->addInitialSource(mockSource.get());
+
+    auto next = pipeline->getNext();
+    ASSERT_BSONOBJ_EQ(next.value().toBson(),
+                      BSON("_id.a" << 10 << "b" << 5 << "_id.b" << -6548868637522515075LL));
+    next = pipeline->getNext();
+    ASSERT_BSONOBJ_EQ(next.value().toBson(),
+                      BSON("_id.a" << 20 << "b" << 1 << "_id.b" << 2598032665634823220LL));
+    ASSERT(!pipeline->getNext());
+}
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/s/config/initial_split_policy.cpp b/src/mongo/db/s/config/initial_split_policy.cpp
index 5bdfb55d6c3..85b8f58cc90 100644
--- a/src/mongo/db/s/config/initial_split_policy.cpp
+++ b/src/mongo/db/s/config/initial_split_policy.cpp
@@ -43,6 +43,7 @@
 #include "mongo/db/s/balancer/balancer_policy.h"
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/db/vector_clock.h"
+#include "mongo/logv2/log.h"
 #include "mongo/s/balancer_configuration.h"
 #include "mongo/s/catalog/type_shard.h"
 #include "mongo/s/grid.h"
@@ -650,34 +651,29 @@ std::vector<BSONObj> ReshardingSplitPolicy::createRawPipeline(const ShardKeyPatt
 
     std::vector<BSONObj> res;
     const auto& shardKeyFields = shardKey.getKeyPatternFields();
-
-    BSONObjBuilder projectValBuilder;
     BSONObjBuilder sortValBuilder;
-
+    using Doc = Document;
+    using Arr = std::vector<Value>;
+    using V = Value;
+    Arr arrayToObjectBuilder;
     for (auto&& fieldRef : shardKeyFields) {
         // If the shard key includes a hashed field and current fieldRef is the hashed field.
         if (shardKey.isHashedPattern() &&
             fieldRef->dottedField().compare(shardKey.getHashedField().fieldNameStringData()) == 0) {
-            projectValBuilder.append(fieldRef->dottedField(),
-                                     BSON("$toHashedIndexKey"
-                                          << "$" + fieldRef->dottedField()));
+            arrayToObjectBuilder.emplace_back(
+                Doc{{"k", V{fieldRef->dottedField()}},
+                    {"v", Doc{{"$toHashedIndexKey", V{"$" + fieldRef->dottedField()}}}}});
         } else {
-            projectValBuilder.append(
-                str::stream() << fieldRef->dottedField(),
-                BSON("$ifNull" << BSON_ARRAY("$" + fieldRef->dottedField() << BSONNULL)));
+            arrayToObjectBuilder.emplace_back(Doc{
+                {"k", V{fieldRef->dottedField()}},
+                {"v", Doc{{"$ifNull", V{Arr{V{"$" + fieldRef->dottedField()}, V{BSONNULL}}}}}}});
         }
-
         sortValBuilder.append(fieldRef->dottedField().toString(), 1);
     }
-
-    // Do not project _id if it's not part of the shard key.
-    if (!shardKey.hasId()) {
-        projectValBuilder.append("_id", 0);
-    }
-
     res.push_back(BSON("$sample" << BSON("size" << numSplitPoints * samplesPerChunk)));
-    res.push_back(BSON("$project" << projectValBuilder.obj()));
     res.push_back(BSON("$sort" << sortValBuilder.obj()));
+    res.push_back(
+        Doc{{"$replaceWith", Doc{{"$arrayToObject", Arr{V{arrayToObjectBuilder}}}}}}.toBson());
     return res;
 }
 
@@ -800,14 +796,10 @@ void ReshardingSplitPolicy::_appendSplitPointsFromSample(BSONObjSet* splitPoints
 
     while (nextKey && nRemaining > 0) {
         // if key is hashed, nextKey values are already hashed
-        auto result = splitPoints->insert(
-            dotted_path_support::extractElementsBasedOnTemplate(*nextKey, shardKey.toBSON())
-                .getOwned());
-
+        auto result = splitPoints->insert(nextKey->getOwned());
         if (result.second) {
             nRemaining--;
         }
-
         nextKey = _samples->getNext();
     }
 }
@@ -819,7 +811,6 @@ ReshardingSplitPolicy::_makePipelineDocumentSource(OperationContext* opCtx,
                                                    int numInitialChunks,
                                                    int samplesPerChunk) {
     auto rawPipeline = createRawPipeline(shardKey, numInitialChunks - 1, samplesPerChunk);
-
     StringMap<ExpressionContext::ResolvedNamespace> resolvedNamespaces;
     resolvedNamespaces[ns.coll()] = {ns, std::vector<BSONObj>{}};
 
diff --git a/src/mongo/db/s/config/initial_split_policy_test.cpp b/src/mongo/db/s/config/initial_split_policy_test.cpp
index 75d1f2ae0a1..d1de8e1926b 100644
--- a/src/mongo/db/s/config/initial_split_policy_test.cpp
+++ b/src/mongo/db/s/config/initial_split_policy_test.cpp
@@ -35,6 +35,7 @@
 #include "mongo/db/s/config/config_server_test_fixture.h"
 #include "mongo/db/s/config/initial_split_policy.h"
 #include "mongo/db/vector_clock.h"
+#include "mongo/logv2/log.h"
 #include "mongo/s/catalog/type_shard.h"
 #include "mongo/s/catalog/type_tags.h"
 #include "mongo/unittest/unittest.h"
@@ -1765,9 +1766,9 @@ TEST_F(ReshardingInitSplitTest, NoZones) {
     shardRegistry()->reload(operationContext());
 
     std::list<BSONObj> mockSamples;
-    mockSamples.push_back(BSON("x" << 10 << "y" << 10));
-    mockSamples.push_back(BSON("x" << 10 << "y" << 20));
-    mockSamples.push_back(BSON("x" << 10 << "y" << 30));
+    mockSamples.push_back(BSON("y" << 10));
+    mockSamples.push_back(BSON("y" << 20));
+    mockSamples.push_back(BSON("y" << 30));
 
     auto mockSampleSource = std::make_unique<MockPipelineSource>(std::move(mockSamples));
 
@@ -1802,9 +1803,9 @@ TEST_F(ReshardingInitSplitTest, HashedShardKey) {
     shardRegistry()->reload(operationContext());
 
     std::list<BSONObj> mockSamples;
-    mockSamples.push_back(BSON("x" << 10 << "y" << 7766103514953448109LL));
-    mockSamples.push_back(BSON("x" << 10 << "y" << -9117533237618642180LL));
-    mockSamples.push_back(BSON("x" << 10 << "y" << -1196399207910989725LL));
+    mockSamples.push_back(BSON("y" << 7766103514953448109LL));
+    mockSamples.push_back(BSON("y" << -9117533237618642180LL));
+    mockSamples.push_back(BSON("y" << -1196399207910989725LL));
 
     auto mockSampleSource = std::make_unique<MockPipelineSource>(std::move(mockSamples));
 
@@ -1867,9 +1868,9 @@ TEST_F(ReshardingInitSplitTest, ZonesCoversEntireDomainButInsufficient) {
     shardRegistry()->reload(operationContext());
 
     std::list<BSONObj> mockSamples;
-    mockSamples.push_back(BSON("x" << 10 << "y" << 10));
-    mockSamples.push_back(BSON("x" << 10 << "y" << 20));
-    mockSamples.push_back(BSON("x" << 10 << "y" << 30));
+    mockSamples.push_back(BSON("y" << 10));
+    mockSamples.push_back(BSON("y" << 20));
+    mockSamples.push_back(BSON("y" << 30));
 
     auto mockSampleSource = std::make_unique<MockPipelineSource>(std::move(mockSamples));
 
@@ -1907,9 +1908,9 @@ TEST_F(ReshardingInitSplitTest, SamplesCoincidingWithZones) {
     shardRegistry()->reload(operationContext());
 
     std::list<BSONObj> mockSamples;
-    mockSamples.push_back(BSON("x" << 10 << "y" << 10));
-    mockSamples.push_back(BSON("x" << 10 << "y" << 20));
-    mockSamples.push_back(BSON("x" << 10 << "y" << 30));
+    mockSamples.push_back(BSON("y" << 10));
+    mockSamples.push_back(BSON("y" << 20));
+    mockSamples.push_back(BSON("y" << 30));
 
     auto mockSampleSource = std::make_unique<MockPipelineSource>(std::move(mockSamples));