SERVER-60893 Deduplicate field lookups for a same field path when $group is pushed down to SBE

1 files changed, 103 insertions, 1 deletions

diff --git a/jstests/noPassthroughWithMongod/group_pushdown.js b/jstests/noPassthroughWithMongod/group_pushdown.js
index 13be21a3244..01e355eb4cd 100644
--- a/jstests/noPassthroughWithMongod/group_pushdown.js
+++ b/jstests/noPassthroughWithMongod/group_pushdown.js
@@ -27,7 +27,8 @@ assert.commandWorked(coll.insert([
 
 let assertGroupPushdown = function(coll, pipeline, expectedResults, expectedGroupCountInExplain) {
     const explain = coll.explain().aggregate(pipeline);
-    // When $group isnever pushed down it be present as a stage in the 'winningPlan' of $cursor.
+    // When $group is pushed down it will never be present as a stage in the 'winningPlan' of
+    // $cursor.
     assert.eq(expectedGroupCountInExplain, getAggPlanStages(explain, "GROUP").length, explain);
 
     let results = coll.aggregate(pipeline).toArray();
@@ -98,6 +99,107 @@ assertResultsMatchWithAndWithoutPushdown(
     [{_id: "a", ss: 22}, {_id: "b", ss: 41}, {_id: "c", ss: 15}],
     2);
 
+// The $group stage refers to the same top-level field twice.
+assertResultsMatchWithAndWithoutPushdown(
+    coll,
+    [{$group: {_id: "$item", ps1: {$sum: "$price"}, ps2: {$sum: "$price"}}}],
+    [{_id: "a", ps1: 15, ps2: 15}, {_id: "b", ps1: 30, ps2: 30}, {_id: "c", ps1: 5, ps2: 5}],
+    1);
+
+// The $group stage refers to the same top-level field twice and another top-level field.
+assertResultsMatchWithAndWithoutPushdown(
+    coll,
+    [{
+        $group:
+            {_id: "$item", ps1: {$sum: "$price"}, ps2: {$sum: "$price"}, qs: {$sum: "$quantity"}}
+    }],
+    [
+        {_id: "a", ps1: 15, ps2: 15, qs: 7},
+        {_id: "b", ps1: 30, ps2: 30, qs: 11},
+        {_id: "c", ps1: 5, ps2: 5, qs: 10}
+    ],
+    1);
+
+// The $group stage refers to two existing sub-fields.
+assertResultsMatchWithAndWithoutPushdown(
+    coll,
+    [
+        {$project: {item: 1, price: 1, quantity: 1, dateParts: {$dateToParts: {date: "$date"}}}},
+        {
+            $group: {
+                _id: "$item",
+                hs: {$sum: {$add: ["$dateParts.hour", "$dateParts.hour", "$dateParts.minute"]}}
+            }
+        },
+    ],
+    [{"_id": "a", "hs": 39}, {"_id": "b", "hs": 34}, {"_id": "c", "hs": 23}],
+    1);
+
+// The $group stage refers to a non-existing sub-field twice.
+assertResultsMatchWithAndWithoutPushdown(
+    coll,
+    [{$group: {_id: "$item", hs: {$sum: {$add: ["$date.hour", "$date.hour"]}}}}],
+    [{"_id": "a", "hs": 0}, {"_id": "b", "hs": 0}, {"_id": "c", "hs": 0}],
+    1);
+
+// Two group stages both get pushed down and the second $group stage refers to only existing
+// top-level fields of the first $group. The field name may be one of "result" / "recordId" /
+// "returnKey" / "snapshotId" / "indexId" / "indexKey" / "indexKeyPattern" which are reserved names
+// inside the SBE stage builder. These special names must not hide user-defined field names.
+[[
+    {$group: {_id: "$item", psum: {$sum: "$price"}}},
+    {$group: {_id: "$_id", ss: {$sum: {$add: ["$psum", "$psum"]}}}}
+],
+ [
+     {$group: {_id: "$item", result: {$sum: "$price"}}},
+     {$group: {_id: "$_id", ss: {$sum: {$add: ["$result", "$result"]}}}}
+ ],
+ [
+     {$group: {_id: "$item", recordId: {$sum: "$price"}}},
+     {$group: {_id: "$_id", ss: {$sum: {$add: ["$recordId", "$recordId"]}}}}
+ ],
+ [
+     {$group: {_id: "$item", returnKey: {$sum: "$price"}}},
+     {$group: {_id: "$_id", ss: {$sum: {$add: ["$returnKey", "$returnKey"]}}}}
+ ],
+ [
+     {$group: {_id: "$item", snapshotId: {$sum: "$price"}}},
+     {$group: {_id: "$_id", ss: {$sum: {$add: ["$snapshotId", "$snapshotId"]}}}}
+ ],
+ [
+     {$group: {_id: "$item", indexId: {$sum: "$price"}}},
+     {$group: {_id: "$_id", ss: {$sum: {$add: ["$indexId", "$indexId"]}}}}
+ ],
+ [
+     {$group: {_id: "$item", indexKey: {$sum: "$price"}}},
+     {$group: {_id: "$_id", ss: {$sum: {$add: ["$indexKey", "$indexKey"]}}}}
+ ],
+ [
+     {$group: {_id: "$item", indexKeyPattern: {$sum: "$price"}}},
+     {$group: {_id: "$_id", ss: {$sum: {$add: ["$indexKeyPattern", "$indexKeyPattern"]}}}}
+ ],
+].forEach(pipeline =>
+              assertResultsMatchWithAndWithoutPushdown(
+                  coll, pipeline, [{_id: "a", ss: 30}, {_id: "b", ss: 60}, {_id: "c", ss: 10}], 2));
+
+// The second $group stage refers to both a top-level field and a sub-field twice.
+assertResultsMatchWithAndWithoutPushdown(
+    coll,
+    [
+        {$group: {_id: "$item", ps: {$sum: "$price"}}},
+        {$group: {_id: "$_id", s1: {$sum: "$ps"}, s2: {$sum: {$add: ["$p.a", "$p.a"]}}}}
+    ],
+    [
+        {"_id": "a", "s1": 15, "s2": 0},
+        {"_id": "b", "s1": 30, "s2": 0},
+        {"_id": "c", "s1": 5, "s2": 0}
+    ],
+    2);
+
+// TODO SERVER-59951: Add more test cases that the second $group stage refers to sub-fields when we
+// enable $mergeObject or document id expression. As of now we don't have a way to produce valid
+// subdocuments from a $group stage.
+
 // Run a group with an unsupported accumultor and check that it doesn't get pushed down.
 assertNoGroupPushdown(coll, [{$group: {_id: "$item", s: {$stdDevSamp: "$quantity"}}}], [
     {"_id": "a", "s": 2.1213203435596424},