SERVER-74245 Encode 'ExpressionContext::needsMerge' in SBE plan cache keys

10 files changed, 183 insertions, 43 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index 6bb95c612fc..c7616038666 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -346,6 +346,8 @@ last-continuous:
     ticket: SERVER-74461
   - test_file: src/mongo/db/modules/enterprise/jstests/fle2/top_command.js
     ticket: SERVER-74460
+  - test_file: jstests/sharding/query/group_plan_cache_sharded.js
+    ticket: SERVER-74245
   suites: null
 last-lts:
   all:
diff --git a/jstests/sharding/query/group_plan_cache_sharded.js b/jstests/sharding/query/group_plan_cache_sharded.js
new file mode 100644
index 00000000000..187c4e6cb73
--- /dev/null
+++ b/jstests/sharding/query/group_plan_cache_sharded.js
@@ -0,0 +1,98 @@
+/**
+ * Test which verifies that the SBE plan cache correctly caches different plans for $group depending
+ * on whether the results will be fed into a merging pipeline or not.
+ *
+ * @tags: [
+ *   requires_sharding,
+ *   # The SBE plan cache was enabled by default in 6.3.
+ *   requires_fcv_63,
+ *   # This test uses the _id index
+ *   expects_explicit_underscore_id_index,
+ * ]
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/sbe_util.js");  // For 'checkSBEEnabled'.
+
+const st = new ShardingTest({shards: 2, rs: {nodes: 1}});
+const mongosDB = st.s.getDB(jsTestName());
+
+// This test is specifically verifying the behavior of the SBE plan cache, which is only enabled
+// when SBE is enabled.
+if (!checkSBEEnabled(mongosDB)) {
+    jsTestLog("Skipping test because SBE is not enabled");
+    st.stop();
+    return;
+}
+
+const collName = jsTestName();
+const coll = mongosDB[collName];
+
+function runPipeline(predicateValue) {
+    return coll
+        .aggregate([
+            {$match: {_id: {$gte: predicateValue}}},
+            {$group: {_id: null, sumResult: {$sum: "$a"}}}
+        ])
+        .toArray();
+}
+
+// Shard coll on _id.
+st.shardColl(
+    coll, {_id: 1} /* key */, {_id: 0} /* split */, {_id: 0} /* move */, mongosDB.getName());
+const docs = [{_id: -2, a: 1}, {_id: 2, a: 2}];
+assert.commandWorked(coll.insertMany(docs));
+
+assert.eq(0, coll.getPlanCache().list().length, "Expected 0 cache entries");
+
+// Run the first aggregate, which will match the second document and target a single shard.
+let res = runPipeline(0);
+assert.eq(res.length, 1);
+assert.eq(res[0], {_id: null, sumResult: 2}, res);
+
+let cacheEntries = coll.getPlanCache().list();
+assert.eq(1, cacheEntries.length, cacheEntries);
+
+// Capture the plan cache key from our lone cache entry.
+const nonMergingCacheEntry = cacheEntries[0];
+assert(nonMergingCacheEntry.hasOwnProperty("planCacheKey"));
+const nonMergingCacheKey = nonMergingCacheEntry.planCacheKey;
+
+// Run the second aggregate, which will match both documents and target both shards.
+res = runPipeline(-2000);
+assert.eq(res.length, 1);
+assert.eq(res[0], {_id: null, sumResult: 3}, res);
+
+// The second aggregate will produce two additional cache entries; one for each shard.
+cacheEntries = coll.getPlanCache().list();
+assert.eq(3, cacheEntries.length, cacheEntries);
+
+// Verify that all three cache entries are pinned and active. Then, verify that exactly one
+// cache entry has the 'nonMergingCacheKey', while the other two have a different key.
+let mergingCacheKey = null;
+let mergingKeyCount = 0;
+let nonMergingKeyCount = 0;
+for (const cacheEntry of cacheEntries) {
+    assert(cacheEntry.isPinned, cacheEntry);
+    assert(cacheEntry.isActive, cacheEntry);
+    assert(cacheEntry.hasOwnProperty("planCacheKey"));
+    if (cacheEntry.planCacheKey === nonMergingCacheKey) {
+        nonMergingKeyCount++;
+    } else {
+        mergingKeyCount++;
+        // If we haven't seen the merging cache key before, stash it so that we can verify that the
+        // two merging plans have the exact same cache key.
+        if (mergingCacheKey === null) {
+            mergingCacheKey = cacheEntry.planCacheKey;
+        } else {
+            assert.eq(cacheEntry.planCacheKey, mergingCacheKey, tojson(cacheEntries));
+        }
+    }
+}
+
+assert.eq(nonMergingKeyCount, 1, tojson(cacheEntries));
+assert.eq(mergingKeyCount, 2, tojson(cacheEntries));
+
+st.stop();
+}());
diff --git a/jstests/sharding/query/merge_use_cases.js b/jstests/sharding/query/merge_use_cases.js
index 80437a0e560..6c1c71b9419 100644
--- a/jstests/sharding/query/merge_use_cases.js
+++ b/jstests/sharding/query/merge_use_cases.js
@@ -107,8 +107,7 @@ runAggregate({startDate: hourSix, whenMatchedMode: "fail", whenNotMatchedMode: "
 
 res = rollupColl.find().sort({_id: 1}).toArray();
 assert.eq(3, res.length, tojson(res));
-// TODO SERVER-74245 Re-enable this assertion once the error has been investigated.
-// assert.eq(res[2], {_id: "2018-08-15T06", ticks: ticksSum, avgTemp: tempSum / samplesPerHour});
+assert.eq(res[2], {_id: "2018-08-15T06", ticks: ticksSum, avgTemp: tempSum / samplesPerHour});
 
 st.stop();
 }());
diff --git a/src/mongo/db/query/canonical_query_encoder.cpp b/src/mongo/db/query/canonical_query_encoder.cpp
index 7173de91f0a..1c89a53442d 100644
--- a/src/mongo/db/query/canonical_query_encoder.cpp
+++ b/src/mongo/db/query/canonical_query_encoder.cpp
@@ -1099,6 +1099,12 @@ std::string encodeSBE(const CanonicalQuery& cq) {
         cq.getOpCtx() && APIParameters::get(cq.getOpCtx()).getAPIStrict().value_or(false);
     bufBuilder.appendChar(apiStrict ? 1 : 0);
 
+    // We can wind up with different query plans for aggregate commands if 'needsMerge' is set or
+    // not. For instance, when 'needsMerge' is true, $group queries will produce partial aggregates
+    // as output, and complete output otherwise.
+    const bool needsMerge = cq.getExpCtx()->needsMerge;
+    bufBuilder.appendChar(needsMerge ? 1 : 0);
+
     encodeFindCommandRequest(cq.getFindCommandRequest(), &bufBuilder);
 
     encodePipeline(cq.pipeline(), &bufBuilder);
diff --git a/src/mongo/db/query/canonical_query_encoder_test.cpp b/src/mongo/db/query/canonical_query_encoder_test.cpp
index 0597bd52eb0..2e554419dca 100644
--- a/src/mongo/db/query/canonical_query_encoder_test.cpp
+++ b/src/mongo/db/query/canonical_query_encoder_test.cpp
@@ -74,7 +74,8 @@ protected:
         BSONObj collation,
         std::unique_ptr<FindCommandRequest> findCommand = nullptr,
         std::vector<BSONObj> pipelineObj = {},
-        bool isCountLike = false) {
+        bool isCountLike = false,
+        bool needsMerge = false) {
         if (!findCommand) {
             findCommand = std::make_unique<FindCommandRequest>(nss);
         }
@@ -85,6 +86,7 @@ protected:
 
         const auto expCtx = make_intrusive<ExpressionContextForTest>(opCtx, nss);
         expCtx->addResolvedNamespaces({foreignNss});
+        expCtx->needsMerge = needsMerge;
         if (!findCommand->getCollation().isEmpty()) {
             auto statusWithCollator = CollatorFactoryInterface::get(opCtx->getServiceContext())
                                           ->makeFromBSON(findCommand->getCollation());
@@ -148,7 +150,8 @@ protected:
                            const char* projStr,
                            std::unique_ptr<FindCommandRequest> findCommand = nullptr,
                            std::vector<BSONObj> pipelineObj = {},
-                           bool isCountLike = false) {
+                           bool isCountLike = false,
+                           bool needsMerge = false) {
         auto& stream = gctx.outStream();
         stream << "==== VARIATION: sbe, query=" << queryStr << ", sort=" << sortStr
                << ", proj=" << projStr;
@@ -161,6 +164,9 @@ protected:
         if (isCountLike) {
             stream << ", isCountLike=true";
         }
+        if (needsMerge) {
+            stream << ", needsMerge=true";
+        }
         stream << std::endl;
         BSONObj collation;
         unique_ptr<CanonicalQuery> cq(canonicalize(opCtx(),
@@ -170,7 +176,8 @@ protected:
                                                    collation,
                                                    std::move(findCommand),
                                                    std::move(pipelineObj),
-                                                   isCountLike));
+                                                   isCountLike,
+                                                   needsMerge));
         cq->setSbeCompatible(true);
         const auto key = canonical_query_encoder::encodeSBE(*cq);
         gctx.outStream() << key << std::endl;
@@ -581,5 +588,29 @@ TEST_F(CanonicalQueryEncoderTest, ComputeKeyWithApiStrict) {
     }
 }
 
+TEST_F(CanonicalQueryEncoderTest, ComputeKeyWithNeedsMerge) {
+    unittest::GoldenTestContext gctx(&goldenTestConfig);
+    RAIIServerParameterControllerForTest controllerSBE("internalQueryFrameworkControl",
+                                                       "trySbeEngine");
+    const auto groupStage = fromjson("{$group: {_id: '$a', out: {$sum: 1}}}");
+    testComputeSBEKey(gctx,
+                      "{}",
+                      "{}",
+                      "{}",
+                      nullptr /* findCommand */,
+                      {groupStage},
+                      false /* isCountLike */,
+                      false /* needsMerge */);
+
+    testComputeSBEKey(gctx,
+                      "{}",
+                      "{}",
+                      "{}",
+                      nullptr /* findCommand */,
+                      {groupStage},
+                      false /* isCountLike */,
+                      true /* needsMerge */);
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e.txt b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e.txt
index a7e6f5d8c82..d5f6f864cce 100644
--- a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e.txt
+++ b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e.txt
@@ -1,54 +1,54 @@
 ==== VARIATION: sbe, query={}, sort={}, proj={}
-YW4ABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+YW4ABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={$or: [{a: 1}, {b: 2}]}, sort={}, proj={}
-b3IAW2VxAGE/AAAAACxlcQBiPwEAAABdBQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+b3IAW2VxAGE/AAAAACxlcQBiPwEAAABdBQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}
-ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={b: 1}, sort={}, proj={}
-ZXEAYj8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYj8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={a: 1, b: 1, c: 1}, sort={}, proj={}
-YW4AW2VxAGE/AAAAACxlcQBiPwEAAAAsZXEAYz8CAAAAXQUAAAAAfHx8fAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
+YW4AW2VxAGE/AAAAACxlcQBiPwEAAAAsZXEAYz8CAAAAXQUAAAAAfHx8fAAAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
 ==== VARIATION: sbe, query={}, sort={a: 1}, proj={}
-YW4ABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+YW4ABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={}, sort={a: -1}, proj={}
-YW4ABQAAAAB8ZGF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+YW4ABQAAAAB8ZGF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={a: 1}
-ZXEAYT8AAAAADAAAABBhAAEAAAAAfGFhfHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYT8AAAAADAAAABBhAAEAAAAAfGFhfHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={}, sort={a: 1}, proj={a: 1}
-YW4ADAAAABBhAAEAAAAAfGFhfHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+YW4ADAAAABBhAAEAAAAAfGFhfHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={}, sort={a: 1}, proj={a: 1}
-YW4ADAAAABBhAAEAAAAAfGFhfHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+YW4ADAAAABBhAAEAAAAAfGFhfHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={}, sort={}, proj={a: 1}
-YW4ADAAAABBhAAEAAAAAfHx8fAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
+YW4ADAAAABBhAAEAAAAAfHx8fAAAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
 ==== VARIATION: sbe, query={}, sort={}, proj={a: true}
-YW4ACQAAAAhhAAEAfHx8fAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
+YW4ACQAAAAhhAAEAfHx8fAAAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
 ==== VARIATION: sbe, query={}, sort={}, proj={a: false}
-YW4ACQAAAAhhAAAAfHx8fAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
+YW4ACQAAAAhhAAAAfHx8fAAAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
 ==== VARIATION: sbe, query={}, sort={}, proj={}, isCountLike=true
-YW4ABQAAAAB8fHx8AAEAAAAAAAAAAABubm5uBQAAAABmfA==
+YW4ABQAAAAB8fHx8AAEAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=1, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAHRubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAAB0bm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAGZubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABmbm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=1, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG50bm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABudG5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5uZm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubmZuBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAKAAAAAAAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAACgAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAACgAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAoAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}, hint={ $natural: 1 }, allowDiskUse=0, returnKey=0, requestResumeToken=1
-ZXEAYT8AAAAABQAAAAB8fHwTAAAAECRuYXR1cmFsAAEAAAAAfAAAAAAAAAAAAAAAbm50bhgAAAASJHJlY29yZElkAAEAAAAAAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8fHwTAAAAECRuYXR1cmFsAAEAAAAAfAAAAAAAAAAAAAAAAG5udG4YAAAAEiRyZWNvcmRJZAABAAAAAAAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}, hint={ a: 1 }, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8fHwMAAAAEGEAAQAAAAB8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYT8AAAAABQAAAAB8fHwMAAAAEGEAAQAAAAB8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}, hint={ a: -1 }, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8fHwMAAAAEGEA/////wB8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYT8AAAAABQAAAAB8fHwMAAAAEGEA/////wB8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
diff --git a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_pipeline.txt b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_pipeline.txt
index 608347f42ca..97be166707c 100644
--- a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_pipeline.txt
+++ b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_pipeline.txt
@@ -1,12 +1,12 @@
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}
-ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}
-ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfFoAAAADJGxvb2t1cABMAAAAAmZyb20ADAAAAGZvcmVpZ25jb2xsAAJhcwADAAAAYXMAAmxvY2FsRmllbGQAAgAAAGEAAmZvcmVpZ25GaWVsZAACAAAAYgAAAA==
+ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnxaAAAAAyRsb29rdXAATAAAAAJmcm9tAAwAAABmb3JlaWduY29sbAACYXMAAwAAAGFzAAJsb2NhbEZpZWxkAAIAAABhAAJmb3JlaWduRmllbGQAAgAAAGIAAAA=
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}
-ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfFsAAAADJGxvb2t1cABNAAAAAmZyb20ADAAAAGZvcmVpZ25jb2xsAAJhcwADAAAAYXMAAmxvY2FsRmllbGQAAwAAAGExAAJmb3JlaWduRmllbGQAAgAAAGIAAAA=
+ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnxbAAAAAyRsb29rdXAATQAAAAJmcm9tAAwAAABmb3JlaWduY29sbAACYXMAAwAAAGFzAAJsb2NhbEZpZWxkAAMAAABhMQACZm9yZWlnbkZpZWxkAAIAAABiAAAA
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}
-ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfFsAAAADJGxvb2t1cABNAAAAAmZyb20ADAAAAGZvcmVpZ25jb2xsAAJhcwADAAAAYXMAAmxvY2FsRmllbGQAAgAAAGEAAmZvcmVpZ25GaWVsZAADAAAAYjEAAAA=
+ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnxbAAAAAyRsb29rdXAATQAAAAJmcm9tAAwAAABmb3JlaWduY29sbAACYXMAAwAAAGFzAAJsb2NhbEZpZWxkAAIAAABhAAJmb3JlaWduRmllbGQAAwAAAGIxAAAA
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}
-ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfFsAAAADJGxvb2t1cABNAAAAAmZyb20ADAAAAGZvcmVpZ25jb2xsAAJhcwAEAAAAYXMxAAJsb2NhbEZpZWxkAAIAAABhAAJmb3JlaWduRmllbGQAAgAAAGIAAAA=
+ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnxbAAAAAyRsb29rdXAATQAAAAJmcm9tAAwAAABmb3JlaWduY29sbAACYXMABAAAAGFzMQACbG9jYWxGaWVsZAACAAAAYQACZm9yZWlnbkZpZWxkAAIAAABiAAAA
 ==== VARIATION: sbe, query={a: 1}, sort={}, proj={}
-ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfFoAAAADJGxvb2t1cABMAAAAAmZyb20ADAAAAGZvcmVpZ25jb2xsAAJhcwADAAAAYXMAAmxvY2FsRmllbGQAAgAAAGEAAmZvcmVpZ25GaWVsZAACAAAAYgAAAF0AAAADJGxvb2t1cABPAAAAAmZyb20ADAAAAGZvcmVpZ25jb2xsAAJhcwAEAAAAYXMxAAJsb2NhbEZpZWxkAAMAAABhMQACZm9yZWlnbkZpZWxkAAMAAABiMQAAAA==
+ZXEAYT8AAAAABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnxaAAAAAyRsb29rdXAATAAAAAJmcm9tAAwAAABmb3JlaWduY29sbAACYXMAAwAAAGFzAAJsb2NhbEZpZWxkAAIAAABhAAJmb3JlaWduRmllbGQAAgAAAGIAAABdAAAAAyRsb29rdXAATwAAAAJmcm9tAAwAAABmb3JlaWduY29sbAACYXMABAAAAGFzMQACbG9jYWxGaWVsZAADAAAAYTEAAmZvcmVpZ25GaWVsZAADAAAAYjEAAAA=
diff --git a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_read_concern.txt b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_read_concern.txt
index 5b1cbb1651f..3583bcc1485 100644
--- a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_read_concern.txt
+++ b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_s_b_e_with_read_concern.txt
@@ -1,6 +1,6 @@
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAGZ8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAABmfA==
 ==== VARIATION: sbe, query={a: 1}, sort={a: 1}, proj={}, hint={}, allowDiskUse=0, returnKey=0, requestResumeToken=0
-ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAG5ubm4FAAAAAHR8
+ZXEAYT8AAAAABQAAAAB8YWF8fHwAAAAAAAAAAAAAAABubm5uBQAAAAB0fA==
diff --git a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_with_api_strict.txt b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_with_api_strict.txt
index f659839252d..150e5b4ddb5 100644
--- a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_with_api_strict.txt
+++ b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_with_api_strict.txt
@@ -3,6 +3,6 @@ an||||ff
 ==== VARIATION: query={}, sort={}, proj={}
 an||||ft
 ==== VARIATION: sbe, query={}, sort={}, proj={}
-YW4ABQAAAAB8fHx8AAAAAAAAAAAAAABubm5uBQAAAABmfA==
+YW4ABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw=
 ==== VARIATION: sbe, query={}, sort={}, proj={}
-YW4ABQAAAAB8fHx8AAABAAAAAAAAAABubm5uBQAAAABmfA==
+YW4ABQAAAAB8fHx8AAABAAAAAAAAAAAAbm5ubgUAAAAAZnw=
diff --git a/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_with_needs_merge.txt b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_with_needs_merge.txt
new file mode 100644
index 00000000000..6e10154b4dd
--- /dev/null
+++ b/src/mongo/db/test_output/query/canonical_query_encoder_test/compute_key_with_needs_merge.txt
@@ -0,0 +1,4 @@
+==== VARIATION: sbe, query={}, sort={}, proj={}
+YW4ABQAAAAB8fHx8AAAAAAAAAAAAAAAAbm5ubgUAAAAAZnw/AAAAAyRncm91cAAyAAAAAl9pZAADAAAAJGEAA291dAAcAAAAAyRzdW0AEQAAABAkY29uc3QAAQAAAAAAAAA=
+==== VARIATION: sbe, query={}, sort={}, proj={}, needsMerge=true
+YW4ABQAAAAB8fHx8AAAAAQAAAAAAAAAAbm5ubgUAAAAAZnw/AAAAAyRncm91cAAyAAAAAl9pZAADAAAAJGEAA291dAAcAAAAAyRzdW0AEQAAABAkY29uc3QAAQAAAAAAAAA=