From bc4d929b684303f89db3fcb8e33b097fa646f553 Mon Sep 17 00:00:00 2001
From: Rui Liu <lriuui0x0@gmail.com>
Date: Fri, 16 Dec 2022 11:27:40 +0000
Subject: SERVER-71663 Do not materialize document when group node doesn't have
 dependency

---
 jstests/noPassthrough/read_only_allow_disk_use.js    |  9 +++++++--
 .../noPassthrough/views_count_distinct_disk_use.js   |  8 +++++++-
 src/mongo/db/query/sbe_stage_builder.cpp             | 20 ++++++++++++++++++++
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/jstests/noPassthrough/read_only_allow_disk_use.js b/jstests/noPassthrough/read_only_allow_disk_use.js
index f3be69cbd4b..8162d6c736d 100644
--- a/jstests/noPassthrough/read_only_allow_disk_use.js
+++ b/jstests/noPassthrough/read_only_allow_disk_use.js
@@ -13,6 +13,8 @@
 (function() {
 "use strict";
 
+load("jstests/libs/sbe_util.js");  // For checkSBEEnabled.
+
 const memoryLimitMb = 1;
 const memoryLimitBytes = 1 * 1024 * 1024;
 const largeStr = "A".repeat(1024 * 1024);  // 1MB string
@@ -104,8 +106,11 @@ function runTest(conn, allowDiskUseByDefault) {
     // The 'count' command within the memory limit must pass.
     assert.eq(coll.count(), memoryLimitMb + 1);
 
-    // The 'count' command exceeding the memory limit must fail.
-    assertFailed({count: view.getName()});
+    // In SBE $sort and $count will not cause spilling, because the largeStr is not saved in memory.
+    // Otherwise, the 'count' command exceeding the memory limit must fail.
+    if (!checkSBEEnabled(testDb)) {
+        assertFailed({count: view.getName()});
+    }
 
     // The 'distinct' command within the memory limit must pass.
     assert.eq(coll.distinct("x"), [0, 1]);
diff --git a/jstests/noPassthrough/views_count_distinct_disk_use.js b/jstests/noPassthrough/views_count_distinct_disk_use.js
index 553509bf983..cd8d3046637 100644
--- a/jstests/noPassthrough/views_count_distinct_disk_use.js
+++ b/jstests/noPassthrough/views_count_distinct_disk_use.js
@@ -4,6 +4,8 @@
 (function() {
 "use strict";
 
+load("jstests/libs/sbe_util.js");  // For checkSBEEnabled.
+
 const conn = MongoRunner.runMongod();
 assert.neq(null, conn, "mongod was unable to start up");
 
@@ -33,7 +35,11 @@ function testDiskUse(cmd) {
 // stage needs to spill to disk if the memory limit is reached.
 assert.commandWorked(viewsDB.adminCommand(
     {setParameter: 1, internalQueryMaxBlockingSortMemoryUsageBytes: memoryLimitMb * 1024 * 1024}));
-testDiskUse({count: "largeView"});
+
+// In SBE the $sort will not cause spilling because it's only the integers being sorted on.
+if (!checkSBEEnabled(viewsDB)) {
+    testDiskUse({count: "largeView"});
+}
 
 // The 'distinct' command executes the view definition pipeline containing the '$sort' stage. This
 // stage needs to spill to disk if the memory limit is reached.
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index 098b9b965f2..7fff4840592 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -2540,6 +2540,26 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder
         childReqs.clear(kResult);
     }
 
+    if (!groupNode->needWholeDocument) {
+        // If the group node doesn't have any dependency (e.g. $count) or if the dependency can be
+        // satisfied by the child node (e.g. covered index scan), we can clear the kResult
+        // requirement for the child.
+        if (groupNode->requiredFields.empty()) {
+            childReqs.clear(kResult);
+        } else if (childNode->getType() == StageType::STAGE_PROJECTION_COVERED) {
+            auto pn = static_cast<const ProjectionNodeCovered*>(childNode);
+            std::set<std::string> providedFieldSet;
+            for (auto&& elt : pn->coveredKeyObj) {
+                providedFieldSet.emplace(elt.fieldNameStringData());
+            }
+            if (std::all_of(groupNode->requiredFields.begin(),
+                            groupNode->requiredFields.end(),
+                            [&](const std::string& f) { return providedFieldSet.count(f); })) {
+                childReqs.clear(kResult);
+            }
+        }
+    }
+
     // Builds the child and gets the child result slot.
     auto [childStage, childOutputs] = build(childNode, childReqs);
 
-- 
cgit v1.2.1