SERVER-51843 [SBE] Support MERGE_SORT stage with children that are not fetched

12 files changed, 383 insertions, 121 deletions

diff --git a/buildscripts/resmokeconfig/suites/core_minimum_batch_size.yml b/buildscripts/resmokeconfig/suites/core_minimum_batch_size.yml
index 10a7aa18aa6..3df148643bb 100644
--- a/buildscripts/resmokeconfig/suites/core_minimum_batch_size.yml
+++ b/buildscripts/resmokeconfig/suites/core_minimum_batch_size.yml
@@ -13,6 +13,8 @@ selector:
   - jstests/core/profile2.js # Extra operation for a getmore.
   - jstests/core/sortk.js # Negative limit value changes result to batchSize.
   - jstests/core/tailable_skip_limit.js # Negative limit value changes result to batchSize.
+  # This test fails when SBE is enabled depending on the data generated.
+  - jstests/core/sort_merge.js
 
 executor:
   archive:
diff --git a/jstests/core/explode_for_sort_fetch.js b/jstests/core/explode_for_sort_fetch.js
index cd347efdaca..798ff0e2996 100644
--- a/jstests/core/explode_for_sort_fetch.js
+++ b/jstests/core/explode_for_sort_fetch.js
@@ -4,7 +4,6 @@
  * @tags: [
  *   # Does not work with legacy shellWriteMode.
  *   requires_find_command,
- *   sbe_incompatible,
  * ]
  */
 (function() {
diff --git a/jstests/core/hashed_index_sort.js b/jstests/core/hashed_index_sort.js
index 9c2f3ad5c6d..8bac29774b6 100644
--- a/jstests/core/hashed_index_sort.js
+++ b/jstests/core/hashed_index_sort.js
@@ -7,7 +7,6 @@
  *   # Sort optimizations added in 4.7 can generate a different plan in the presence of equality
  *   # predicates.
  *   requires_fcv_47,
- *   sbe_incompatible,
  * ]
  */
 (function() {
@@ -15,6 +14,12 @@
 
 load("jstests/libs/analyze_plan.js");  // For assertStagesForExplainOfCommand().
 
+// The explain output depends on which execution engine is enabled.
+// TODO Remove when SERVER-51823 lands.
+const isSBEEnabled = (() => {
+    const getParam = db.adminCommand({getParameter: 1, featureFlagSBE: 1});
+    return getParam.hasOwnProperty("featureFlagSBE") && getParam.featureFlagSBE.value;
+})();
 const coll = db.hashed_index_sort;
 coll.drop();
 
@@ -102,7 +107,7 @@ validateFindCmdOutputAndPlan({
         {a: 1, b: 3, c: 4, d: -2},
         {a: 1, b: 4, c: 5, d: -3},
     ],
-    expectedStages: ["IXSCAN", "SORT"]
+    expectedStages: ["IXSCAN", isSBEEnabled ? "SORT_SIMPLE" : "SORT"]
 });
 
 /**
@@ -164,7 +169,7 @@ validateFindCmdOutputAndPlan({
     filter: {},
     project: {_id: 0, a: 1, b: 1},
     sort: {a: 1, b: -1, c: 1},
-    expectedStages: ["SORT", "COLLSCAN"],
+    expectedStages: [isSBEEnabled ? "SORT_SIMPLE" : "SORT", "COLLSCAN"],
 });
 
 // Verify that a list of exact match predicates on range field (prefix) and sort with an immediate
@@ -186,7 +191,7 @@ validateFindCmdOutputAndPlan({
     project: {_id: 0, d: 1, b: 1},
     sort: {d: 1},
     expectedOutput: [{b: 4, d: -3}, {b: 3, d: -2}, {b: 2, d: -1}, {b: 1, d: 0}, {b: 0, d: 1}],
-    expectedStages: ["IXSCAN", "SORT"],
+    expectedStages: ["IXSCAN", isSBEEnabled ? "SORT_DEFAULT" : "SORT"],
     stagesNotExpected: ["COLLSCAN", "FETCH"]
 });
 
@@ -196,6 +201,6 @@ validateFindCmdOutputAndPlan({
     project: {_id: 0, c: 1},
     sort: {c: 1},
     expectedOutput: [{c: 2}],
-    expectedStages: ["IXSCAN", "FETCH", "SORT"]
+    expectedStages: ["IXSCAN", "FETCH", isSBEEnabled ? "SORT_SIMPLE" : "SORT"]
 });
 })();
diff --git a/jstests/core/index_sort_within_multiple_point_ranges.js b/jstests/core/index_sort_within_multiple_point_ranges.js
index d75859286b3..6caaad1cb6f 100644
--- a/jstests/core/index_sort_within_multiple_point_ranges.js
+++ b/jstests/core/index_sort_within_multiple_point_ranges.js
@@ -9,7 +9,6 @@
  * - the query asks for a sort within each point range
  * @tags: [
  *   assumes_no_implicit_collection_creation_after_drop,
- *   sbe_incompatible,
  * ]
  */
 (function() {
diff --git a/jstests/core/merge_sort_collation.js b/jstests/core/merge_sort_collation.js
index 48141d14052..25b8af93c24 100644
--- a/jstests/core/merge_sort_collation.js
+++ b/jstests/core/merge_sort_collation.js
@@ -4,7 +4,6 @@
  *  2) MERGE_SORT stage execution.
  * @tags: [
  *   requires_find_command,
- *   sbe_incompatible,
  * ]
  */
 (function() {
diff --git a/jstests/core/or6.js b/jstests/core/or6.js
index 7b18e1c219a..796be83549e 100644
--- a/jstests/core/or6.js
+++ b/jstests/core/or6.js
@@ -1,7 +1,4 @@
 // A few rooted $or cases.
-// @tags: [
-//   sbe_incompatible,
-// ]
 
 var t = db.jstests_orq;
 t.drop();
diff --git a/jstests/core/sort_merge.js b/jstests/core/sort_merge.js
index 4c1ebab85d0..de59cbd48ba 100644
--- a/jstests/core/sort_merge.js
+++ b/jstests/core/sort_merge.js
@@ -4,11 +4,15 @@
 (function() {
 "use strict";
 
+load("jstests/libs/analyze_plan.js");
+load("jstests/libs/fixture_helpers.js");  // For 'isMongos'.
+
 const coll = db.sort_merge;
 coll.drop();
 
 assert.commandWorked(coll.createIndex({filterFieldA: 1, sortFieldA: 1, sortFieldB: 1}));
 assert.commandWorked(coll.createIndex({filterFieldA: 1, sortFieldA: -1, sortFieldB: -1}));
+assert.commandWorked(coll.createIndex({sortFieldA: 1, sortFieldB: 1}));
 
 // Insert some random data and dump the seed to the log.
 const seed = new Date().getTime();
@@ -39,6 +43,72 @@ function isSorted(array, lessThanFunction) {
     return true;
 }
 
+// Verifies that 'sortMerge' is covered, that is, it has no child FETCH or COLLSCAN stages.
+function verifyCoveredPlan(explain, sortMerge) {
+    assert(isIndexOnly(db, sortMerge), explain);
+}
+
+// Verifies that 'sortMerge' is not a covered plan. In particular, we check that the number of
+// FETCH stages matches the number of branches and that each FETCH has an IXSCAN as a child.
+function verifyNonCoveredPlan(explain, sortMerge) {
+    assert(sortMerge.hasOwnProperty("inputStages"), explain);
+    const numBranches = sortMerge.inputStages.length;
+    const fetchStages = getPlanStages(sortMerge, "FETCH");
+
+    assert.eq(fetchStages.length, numBranches, explain);
+    for (const fetch of fetchStages) {
+        assert(isIxscan(db, fetch), explain);
+    }
+}
+
+// Verifies that 'sortMerge', which is a plan with a mix of non-covered and covered branches,
+// has as many IXSCAN stages as there are branches.
+function verifyMixedPlan(explain, sortMerge) {
+    assert(sortMerge.hasOwnProperty("inputStages"), explain);
+    const numBranches = sortMerge.inputStages.length;
+    const ixscanStages = getPlanStages(sortMerge, "IXSCAN");
+
+    // Note that this is not as strong an assertion as 'isIndexOnly' as some branches will have
+    // FETCH stages.
+    assert.eq(ixscanStages.length, numBranches, explain);
+}
+
+// Verifies that the query solution tree produced by 'sort' and 'filter' produces a SORT_MERGE
+// plan and invokes 'callback' to make custom assertions about the SORT_MERGE plan.
+function verifyPlan(sort, filter, callback) {
+    const explain = coll.find(filter).sort(sort.sortPattern).explain("queryPlanner");
+
+    // Search for all SORT_MERGE stages in case this is a sharded collection.
+    if (FixtureHelpers.isMongos(db)) {
+        const sortMergeStages = getPlanStages(explain, "SORT_MERGE");
+        for (const sortMerge of sortMergeStages) {
+            callback(explain, sortMerge);
+        }
+    } else {
+        const sortMerge = getPlanStage(explain, "SORT_MERGE");
+        callback(explain, sortMerge);
+    }
+}
+
+function runTest(sorts, filters, verifyCallback) {
+    for (let sortInfo of sorts) {
+        for (let filter of filters) {
+            verifyPlan(sortInfo, filter, verifyCallback);
+            const res = coll.find(filter).sort(sortInfo.sortPattern).toArray();
+            assert(isSorted(res, sortInfo.cmpFunction),
+                   () => "Assertion failed for filter: " + filter + "\n" +
+                       "sort pattern " + sortInfo.sortPattern);
+
+            // Check that there are no duplicates.
+            let ids = new Set();
+            for (let doc of res) {
+                assert(!ids.has(doc._id), () => "Duplicate _id: " + tojson(_id));
+                ids.add(doc._id);
+            }
+        }
+    }
+}
+
 // Cases where the children of the $or all require a FETCH.
 (function testFetchedChildren() {
     const kFilterPredicates = [
@@ -77,47 +147,173 @@ function isSorted(array, lessThanFunction) {
         {
             sortPattern: {sortFieldA: 1, sortFieldB: 1},
             cmpFunction: (docA, docB) => docA.sortFieldA < docB.sortFieldA ||
-                (docA.sortFieldA == docB.sortFieldA && docA.sortFieldB < docB.sortFieldB)
+                (docA.sortFieldA === docB.sortFieldA && docA.sortFieldB < docB.sortFieldB)
         },
         {
             sortPattern: {sortFieldA: -1, sortFieldB: -1},
             cmpFunction: (docA, docB) => docA.sortFieldA > docB.sortFieldA ||
-                (docA.sortFieldA == docB.sortFieldA && docA.sortFieldB > docB.sortFieldB)
+                (docA.sortFieldA === docB.sortFieldA && docA.sortFieldB > docB.sortFieldB)
         },
     ];
 
-    for (let sortInfo of kSorts) {
-        for (let filter of kFilterPredicates) {
-            // Check that the results are in order.
-            let res = coll.find(filter).sort(sortInfo.sortPattern).toArray();
-            assert(isSorted(res, sortInfo.cmpFunction),
-                   () => "Assertion failed for filter: " + filter + "\n" +
-                       "sort pattern " + sortInfo.sortPattern);
+    runTest(kSorts, kFilterPredicates, verifyNonCoveredPlan);
+})();
 
-            // Check that there are no duplicates.
-            let ids = new Set();
-            for (let doc of res) {
-                assert(!ids.has(doc._id), () => "Duplicate _id: " + tojson(_id));
-                ids.add(doc._id);
-            }
+// Cases where the children of the $or are NOT fetched.
+(function testUnfetchedChildren() {
+    const kFilterPredicates = [
+        {
+            $or: [
+                {sortFieldA: 1, sortFieldB: 3},
+                {filterFieldA: 2, sortFieldB: 2},
+            ]
+        },
+        {
+            $or: [
+                {sortFieldA: 1, sortFieldB: 3},
+                {sortFieldA: 2, sortFieldB: 2},
+                {sortFieldA: 4, sortFieldB: 2},
+            ]
+        },
+        {
+            $or: [
+                {sortFieldA: 1, sortFieldB: 3},
+                {sortFieldA: 2, sortFieldB: 2},
+                {sortFieldA: 4, sortFieldB: 2},
+                {sortFieldA: 3, sortFieldB: 1},
+            ]
+        },
+    ];
+
+    const kSorts = [
+        {
+            sortPattern: {sortFieldB: 1},
+            cmpFunction: (docA, docB) => docA.sortFieldB < docB.sortFieldB
+        },
+        {
+            sortPattern: {sortFieldA: -1},
+            cmpFunction: (docA, docB) => docA.sortFieldA > docB.sortFieldA
+        },
+        {
+            sortPattern: {sortFieldA: 1, sortFieldB: 1},
+            cmpFunction: (docA, docB) => docA.sortFieldA < docB.sortFieldA ||
+                (docA.sortFieldA === docB.sortFieldA && docA.sortFieldB < docB.sortFieldB)
+        },
+        {
+            sortPattern: {sortFieldB: 1, sortFieldA: 1},
+            cmpFunction: (docA, docB) => docA.sortFieldB < docB.sortFieldB ||
+                (docA.sortFieldB === docB.sortFieldB && docA.sortFieldA < docB.sortFieldA)
+        },
+        {
+            sortPattern: {sortFieldA: 1, sortFieldB: -1},
+            cmpFunction: (docA, docB) => docA.sortFieldA < docB.sortFieldA ||
+                (docA.sortFieldA === docB.sortFieldA && docA.sortFieldA > docB.sortFieldA)
+        },
+    ];
+    runTest(kSorts, kFilterPredicates, verifyCoveredPlan);
+})();
+
+// Cases where the children of the $or are a mix of fetched and unfetched.
+(function testFetchedAndUnfetchedChildren() {
+    const kFilterPredicates = [
+        {$or: [{sortFieldA: 1, sortFieldB: 2}, {sortFieldA: 2, sortFieldB: 2, filterFieldB: 1}]},
+        {
+            $or: [
+                {sortFieldA: 1, sortFieldB: 2},
+                {sortFieldA: 2, sortFieldB: 2, filterFieldB: 1},
+                {sortFieldA: 3, sortFieldB: 4, filterFieldA: 2},
+                {sortFieldA: 4, sortFieldB: 1, filterFieldB: 4}
+            ]
+        },
+    ];
+
+    const kSorts = [
+        {
+            sortPattern: {sortFieldB: 1},
+            cmpFunction: (docA, docB) => docA.sortFieldB < docB.sortFieldB
+        },
+        {
+            sortPattern: {sortFieldA: -1},
+            cmpFunction: (docA, docB) => docA.sortFieldA > docB.sortFieldA
+        },
+        {
+            sortPattern: {sortFieldA: 1, sortFieldB: 1},
+            cmpFunction: (docA, docB) => docA.sortFieldA < docB.sortFieldA ||
+                (docA.sortFieldA === docB.sortFieldA && docA.sortFieldB < docB.sortFieldB)
+        },
+        {
+            sortPattern: {sortFieldB: 1, sortFieldA: 1},
+            cmpFunction: (docA, docB) => docA.sortFieldB < docB.sortFieldB ||
+                (docA.sortFieldB === docB.sortFieldB && docA.sortFieldA < docB.sortFieldA)
+        },
+        {
+            sortPattern: {sortFieldA: 1, sortFieldB: -1},
+            cmpFunction: (docA, docB) => docA.sortFieldA < docB.sortFieldA ||
+                (docA.sortFieldA === docB.sortFieldA && docA.sortFieldA > docB.sortFieldA)
+        },
+    ];
+    runTest(kSorts, kFilterPredicates, verifyMixedPlan);
+})();
+
+// Insert documents with arrays into the collection and check that the deduping works correctly.
+(function testDeduplication() {
+    assert.commandWorked(coll.insert([
+        {filterFieldA: [1, 2], filterFieldB: "multikeydoc", sortFieldA: 1, sortFieldB: 1},
+        {sortFieldA: [1, 2], filterFieldA: "multikeydoc"}
+    ]));
+
+    const kUniqueFilters = [
+        {
+            $or: [
+                {filterFieldA: 1, filterFieldB: "multikeydoc"},
+                {filterFieldA: 2, filterFieldB: "multikeydoc"}
+            ],
+        },
+        {
+            $or: [
+                {sortFieldA: 1, filterFieldA: "multikeydoc"},
+                {sortFieldA: 2, filterFieldA: "multikeydoc"}
+            ]
         }
+    ];
+
+    for (const filter of kUniqueFilters) {
+        // Both branches of the $or will return the same document, which should be de-duped. Make
+        // sure only one document is returned from the server.
+        assert.eq(coll.find(filter).sort({sortFieldA: 1}).itcount(), 1);
     }
 })();
 
-// TODO SERVER-51843: Test with children that are not fetched.
-
-// Insert an arrays into the collection and check that the deduping works correctly.
-assert.commandWorked(
-    coll.insert({filterFieldA: [1, 2], filterFieldB: "multikeydoc", sortFieldA: 1, sortFieldB: 1}));
-// Both branches of the $or will return the same document, which should be de-duped. Make sure
-// only one document is returned from the server.
-assert.eq(coll.find({
-                  $or: [
-                      {filterFieldA: 1, filterFieldB: "multikeydoc"},
-                      {filterFieldA: 2, filterFieldB: "multikeydoc"}
-                  ]
-              })
-              .sort({sortFieldA: 1})
-              .itcount(),
-          1);
+// Verify that sort merge works correctly when sorting dotted paths.
+(function testDottedPathSortMerge() {
+    assert(coll.drop());
+    assert.commandWorked(coll.createIndex({'filterFieldA': 1, 'sortField.a': 1, 'sortField.b': 1}));
+    for (let i = 0; i < 100; ++i) {
+        assert.commandWorked(coll.insert({
+            filterFieldA: randomInt(),
+            filterFieldB: randomInt(),
+            sortField: {a: randomInt(), b: randomInt()}
+        }));
+    }
+
+    const kSortPattern = {
+        sortPattern: {'sortField.a': 1, 'sortField.b': 1},
+        cmpFunction: (docA, docB) => docA.sortField.a < docB.sortField.a ||
+            (docA.sortField.a === docB.sortField.a && docA.sortField.b < docB.sortField.b)
+    };
+
+    const kNonCoveredFilter = {
+        $or: [{filterFieldA: 1, filterFieldB: 2}, {filterFieldA: 2, filterFieldB: 1}]
+    };
+    runTest([kSortPattern], [kNonCoveredFilter], verifyNonCoveredPlan);
+
+    const kCoveredFilter = {
+        $or: [
+            {filterFieldA: 1, 'sortField.a': 1},
+            {filterFieldA: 2, 'sortField.a': 3},
+            {filterFieldA: 3, 'sortField.a': 4, 'sortField.b': 3},
+        ]
+    };
+    runTest([kSortPattern], [kCoveredFilter], verifyCoveredPlan);
+})();
 })();
diff --git a/jstests/core/sort_merge_collation.js b/jstests/core/sort_merge_collation.js
index 8d55360079c..38e6f33d9da 100644
--- a/jstests/core/sort_merge_collation.js
+++ b/jstests/core/sort_merge_collation.js
@@ -8,6 +8,8 @@
 (function() {
 "use strict";
 
+load("jstests/libs/analyze_plan.js");
+
 const numericOrdering = {
     collation: {locale: "en_US", numericOrdering: true}
 };
@@ -21,6 +23,7 @@ assert.commandWorked(
     coll.createIndex({filterFieldA: 1, sortFieldA: 1, sortFieldB: 1}, numericOrdering));
 assert.commandWorked(
     coll.createIndex({filterFieldA: 1, sortFieldA: -1, sortFieldB: -1}, numericOrdering));
+assert.commandWorked(coll.createIndex({sortFieldA: 1, sortFieldB: 1}, numericOrdering));
 
 assert.commandWorked(coll.insert([
     {filterFieldA: "1", filterFieldB: "1", sortFieldA: "18", sortFieldB: "11"},
@@ -54,6 +57,53 @@ function isSorted(array, lessThanFunction) {
     return true;
 }
 
+function runTest(sorts, filters) {
+    for (let sortInfo of sorts) {
+        for (let filter of filters) {
+            // Verify that the sort/filter combination produces a SORT_MERGE plan.
+            const explain = coll.find(filter).sort(sortInfo.sortPattern).explain("queryPlanner");
+            const sortMergeStages = getPlanStages(explain, "SORT_MERGE");
+            assert.gt(sortMergeStages.length, 0, explain);
+
+            // Check that the results are in order.
+            let res = coll.find(filter).sort(sortInfo.sortPattern).toArray();
+            assert(isSorted(res, sortInfo.cmpFunction),
+                   () => "Assertion failed for filter: " + filter + "\n" +
+                       "sort pattern " + sortInfo.sortPattern);
+
+            // Check that there are no duplicates.
+            let ids = new Set();
+            for (let doc of res) {
+                assert(!ids.has(doc._id), () => "Duplicate _id: " + tojson(_id));
+                ids.add(doc._id);
+            }
+        }
+    }
+}
+
+const kSorts = [
+    {
+        sortPattern: {sortFieldA: 1},
+        cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) < parseInt(docB.sortFieldA)
+    },
+    {
+        sortPattern: {sortFieldA: -1},
+        cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) > parseInt(docB.sortFieldA)
+    },
+    {
+        sortPattern: {sortFieldA: 1, sortFieldB: 1},
+        cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) < parseInt(docB.sortFieldA) ||
+            (parseInt(docA.sortFieldA) == parseInt(docB.sortFieldA) &&
+             parseInt(docA.sortFieldB) < parseInt(docB.sortFieldB))
+    },
+    {
+        sortPattern: {sortFieldA: -1, sortFieldB: -1},
+        cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) > parseInt(docB.sortFieldA) ||
+            (parseInt(docA.sortFieldA) == parseInt(docB.sortFieldA) &&
+             parseInt(docA.sortFieldB) > parseInt(docB.sortFieldB))
+    },
+];
+
 // Cases where the children of the $or all require a FETCH.
 (function testFetchedChildren() {
     const kFilterPredicates = [
@@ -80,46 +130,35 @@ function isSorted(array, lessThanFunction) {
         },
     ];
 
-    const kSorts = [
-        {
-            sortPattern: {sortFieldA: 1},
-            cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) < parseInt(docB.sortFieldA)
-        },
-        {
-            sortPattern: {sortFieldA: -1},
-            cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) > parseInt(docB.sortFieldA)
-        },
+    runTest(kSorts, kFilterPredicates);
+})();
+
+// Cases where the children of the $or are IXSCANs.
+(function testUnfetchedChildren() {
+    const kFilterPredicates = [
+        // $or with two children.
+        {$or: [{sortFieldA: "4", sortFieldB: "4"}, {sortFieldA: "3", sortFieldB: "3"}]},
+
+        // $or with three children.
         {
-            sortPattern: {sortFieldA: 1, sortFieldB: 1},
-            cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) < parseInt(docB.sortFieldA) ||
-                (parseInt(docA.sortFieldA) == parseInt(docB.sortFieldA) &&
-                 parseInt(docA.sortFieldB) < parseInt(docB.sortFieldB))
+            $or: [
+                {sortFieldA: "4", sortFieldB: "4"},
+                {sortFieldA: "3", sortFieldB: "3"},
+                {sortFieldA: "7", sortFieldB: "4"}
+            ]
         },
+
+        // $or with four children.
         {
-            sortPattern: {sortFieldA: -1, sortFieldB: -1},
-            cmpFunction: (docA, docB) => parseInt(docA.sortFieldA) > parseInt(docB.sortFieldA) ||
-                (parseInt(docA.sortFieldA) == parseInt(docB.sortFieldA) &&
-                 parseInt(docA.sortFieldB) > parseInt(docB.sortFieldB))
+            $or: [
+                {sortFieldA: "4", sortFieldB: "4"},
+                {sortFieldA: "3", sortFieldB: "3"},
+                {sortFieldA: "7", sortFieldB: "4"},
+                {sortFieldA: "10", sortFieldB: "9"}
+            ]
         },
     ];
 
-    for (let sortInfo of kSorts) {
-        for (let filter of kFilterPredicates) {
-            // Check that the results are in order.
-            let res = coll.find(filter).sort(sortInfo.sortPattern).toArray();
-            assert(isSorted(res, sortInfo.cmpFunction),
-                   () => "Assertion failed for filter: " + filter + "\n" +
-                       "sort pattern " + sortInfo.sortPattern);
-
-            // Check that there are no duplicates.
-            let ids = new Set();
-            for (let doc of res) {
-                assert(!ids.has(doc._id), () => "Duplicate _id: " + tojson(_id));
-                ids.add(doc._id);
-            }
-        }
-    }
+    runTest(kSorts, kFilterPredicates);
 })();
-
-// TODO SERVER-51843: Test with children that are not fetched.
 })();
diff --git a/jstests/core/sorth.js b/jstests/core/sorth.js
index 41563caa205..db1187eb0ec 100644
--- a/jstests/core/sorth.js
+++ b/jstests/core/sorth.js
@@ -1,7 +1,4 @@
 // Tests for the $in/sort/limit optimization combined with inequality bounds.  SERVER-5777
-// @tags: [
-//   sbe_incompatible,
-// ]
 
 (function() {
 "use strict";
diff --git a/jstests/core/sortk.js b/jstests/core/sortk.js
index f1b055c2bcc..8815974cb10 100644
--- a/jstests/core/sortk.js
+++ b/jstests/core/sortk.js
@@ -6,7 +6,6 @@
 //   requires_non_retryable_writes,
 //   # Uses $where operator
 //   requires_scripting,
-//   sbe_incompatible,
 // ]
 
 t = db.jstests_sortk;
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index 9c9312b2c5f..1fdf6fb72fb 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -318,7 +318,6 @@ SlotBasedStageBuilder::makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inpu
 std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder::buildFetch(
     const QuerySolutionNode* root, const PlanStageReqs& reqs) {
     auto fn = static_cast<const FetchNode*>(root);
-    invariant(!reqs.getIndexKeyBitset());
 
     // At present, makeLoopJoinForFetch() doesn't have the necessary logic for producing an
     // oplogTsSlot, so assert that the caller doesn't need oplogTsSlot.
@@ -341,6 +340,11 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     auto relevantSlots = sbe::makeSV();
     outputs.forEachSlot(forwardingReqs, [&](auto&& slot) { relevantSlots.push_back(slot); });
 
+    // Forward slots for components of the index key if our parent requested them.
+    if (auto indexKeySlots = outputs.getIndexKeySlots()) {
+        relevantSlots.insert(relevantSlots.end(), indexKeySlots->begin(), indexKeySlots->end());
+    }
+
     sbe::value::SlotId fetchResultSlot, fetchRecordIdSlot;
     std::tie(fetchResultSlot, fetchRecordIdSlot, stage) = makeLoopJoinForFetch(
         std::move(stage), outputs.get(kRecordId), root->nodeId(), std::move(relevantSlots));
@@ -354,6 +358,10 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
         relevantSlots = sbe::makeSV();
         outputs.forEachSlot(forwardingReqs, [&](auto&& slot) { relevantSlots.push_back(slot); });
 
+        // Forward slots for components of the index key if our parent requested them.
+        if (auto indexKeySlots = outputs.getIndexKeySlots()) {
+            relevantSlots.insert(relevantSlots.end(), indexKeySlots->begin(), indexKeySlots->end());
+        }
 
         std::tie(std::ignore, stage) = generateFilter(_opCtx,
                                                       fn->filter.get(),
@@ -534,19 +542,11 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
 
     auto mergeSortNode = static_cast<const MergeSortNode*>(root);
 
-    uassert(5073803,
-            "SORT_MERGE stage with unfetched children not supported",
-            mergeSortNode->fetched());
-
     const auto sortPattern = SortPattern{mergeSortNode->sort, _cq.getExpCtx()};
     std::vector<sbe::value::SortDirection> direction;
 
     for (const auto& part : sortPattern) {
         uassert(4822881, "Sorting by expression not supported", !part.expression);
-        uassert(4822882,
-                "Sorting by dotted paths not supported",
-                part.fieldPath && part.fieldPath->getPathLength() == 1);
-
         direction.push_back(part.isAscending ? sbe::value::SortDirection::Ascending
                                              : sbe::value::SortDirection::Descending);
     }
@@ -556,46 +556,76 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
     std::vector<sbe::value::SlotVector> inputVals;
 
     // Children must produce all of the slots required by the parent of this SortMergeNode. In
-    // addition to that, children must always produce a 'recordIdSlot' if the 'dedup' flag is true,
-    // and children must always produce a 'resultSlot' because it's needed by the sort logic below.
-    auto childReqs = reqs.copy().set(kResult).setIf(kRecordId, mergeSortNode->dedup);
+    // addition, children must always produce a 'recordIdSlot' if the 'dedup' flag is true.
+    auto childReqs = reqs.copy().setIf(kRecordId, mergeSortNode->dedup);
 
     for (auto&& child : mergeSortNode->children) {
-        auto [stage, outputs] = build(child, childReqs);
-
-        sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projectMap;
         sbe::value::SlotVector inputKeysForChild;
 
-        for (const auto& part : sortPattern) {
-            // Slot holding the sort key.
-            auto sortFieldSlot{_slotIdGenerator.generate()};
-            inputKeysForChild.push_back(sortFieldSlot);
-
-            auto fieldName = part.fieldPath->getFieldName(0);
-            auto fieldNameSV = std::string_view{fieldName.rawData(), fieldName.size()};
-
-            auto getSortFieldExpr = makeFunction("getField"sv,
-                                                 sbe::makeE<sbe::EVariable>(outputs.get(kResult)),
-                                                 sbe::makeE<sbe::EConstant>(fieldNameSV));
-
-            if (auto collatorSlot = _data.env->getSlotIfExists("collator"_sd); collatorSlot) {
-                getSortFieldExpr = makeFunction("collComparisonKey"sv,
-                                                std::move(getSortFieldExpr),
-                                                sbe::makeE<sbe::EVariable>(*collatorSlot));
+        // Map of field name to position within the index key. This is used to account for
+        // mismatches between the sort pattern and the index key pattern. For instance, suppose the
+        // requested sort is {a: 1, b: 1} and the index key pattern is {c: 1, b: 1, a: 1}. When the
+        // slots for the relevant components of the index key are generated (i.e. extract keys for
+        // 'b' and 'a'),  we wish to insert them into 'inputKeys' in the order that they appear in
+        // the sort pattern.
+        StringMap<size_t> indexKeyPositionMap;
+        auto ixnNode = getNodeByType(child, STAGE_IXSCAN);
+        tassert(5184300,
+                str::stream() << "Can't build exec tree for node: " << child->toString(),
+                ixnNode);
+
+        auto ixn = static_cast<const IndexScanNode*>(ixnNode);
+        sbe::IndexKeysInclusionSet indexKeyBitset;
+        size_t i = 0;
+        for (auto&& elt : ixn->index.keyPattern) {
+            for (auto&& sortPart : sortPattern) {
+                auto path = sortPart.fieldPath->fullPath();
+                if (elt.fieldNameStringData() == path) {
+                    indexKeyBitset.set(i);
+                    indexKeyPositionMap.emplace(path, indexKeyPositionMap.size());
+                    break;
+                }
             }
-
-            projectMap.emplace(sortFieldSlot, std::move(getSortFieldExpr));
+            ++i;
         }
+        childReqs.getIndexKeyBitset() = indexKeyBitset;
 
-        stage =
-            sbe::makeS<sbe::ProjectStage>(std::move(stage), std::move(projectMap), root->nodeId());
+        // Children must produce a 'resultSlot' if they produce fetched results.
+        auto [stage, outputs] = build(child, childReqs);
 
-        inputStages.push_back(std::move(stage));
+        tassert(5184301,
+                "SORT_MERGE node must receive a RecordID slot as input from child stage"
+                " if the 'dedup' flag is set",
+                !mergeSortNode->dedup || outputs.has(kRecordId));
 
-        invariant(outputs.has(kResult));
-        invariant(!mergeSortNode->dedup || outputs.has(kRecordId));
+        // Clear the index key bitset after building the child stage.
+        childReqs.getIndexKeyBitset() = boost::none;
+
+        // Insert the index key slots in the order of the sort pattern.
+        auto indexKeys = outputs.extractIndexKeySlots();
+        tassert(5184302,
+                "SORT_MERGE must receive index key slots as input from its child stages",
+                indexKeys.has_value());
+
+        for (const auto& part : sortPattern) {
+            auto partPath = part.fieldPath->fullPath();
+            auto index = indexKeyPositionMap.find(partPath);
+            tassert(5184303,
+                    str::stream() << "Could not find index key position for sort key part "
+                                  << partPath,
+                    index != indexKeyPositionMap.end());
+            auto indexPos = index->second;
+            tassert(5184304,
+                    str::stream() << "Index position " << indexPos
+                                  << " is not less than number of index components "
+                                  << indexKeys->size(),
+                    indexPos < indexKeys->size());
+            auto indexKeyPart = indexKeys->at(indexPos);
+            inputKeysForChild.push_back(indexKeyPart);
+        }
 
         inputKeys.push_back(std::move(inputKeysForChild));
+        inputStages.push_back(std::move(stage));
 
         auto sv = sbe::makeSV();
         outputs.forEachSlot(childReqs, [&](auto&& slot) { sv.push_back(slot); });
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index 4dc17aba683..4a9d277ab00 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -118,8 +118,8 @@ public:
 private:
     StringMap<sbe::value::SlotId> _slots;
 
-    // When an index scan produces parts of an index key for a covered projection, this is where
-    // the slots for the produced values are stored.
+    // When an index scan produces parts of an index key for a covered plan, this is where the
+    // slots for the produced values are stored.
     boost::optional<sbe::value::SlotVector> _indexKeySlots;
 };