SERVER-67284 Allow $fill 'partitionByFields' to partition by a dotted path

2 files changed, 123 insertions, 45 deletions

diff --git a/jstests/aggregation/sources/fill/fill.js b/jstests/aggregation/sources/fill/fill.js
index 874cc2dd55e..332bc7f7e4c 100644
--- a/jstests/aggregation/sources/fill/fill.js
+++ b/jstests/aggregation/sources/fill/fill.js
@@ -17,16 +17,16 @@ load("jstests/aggregation/extras/utils.js");  // For arrayEq.
 const coll = db[jsTestName()];
 coll.drop();
 const documents = [
-    {_id: 1, linear: 1, other: 1, part: 1},
-    {_id: 2, linear: 1, other: 1, part: 2},
-    {_id: 3, linear: null, other: null, part: 1},
-    {_id: 4, linear: null, other: null, part: 2},
-    {_id: 5, linear: 5, other: 10, part: 1},
-    {_id: 6, linear: 6, other: 2, part: 2},
-    {_id: 7, linear: null, other: null, part: 1},
-    {_id: 8, linear: 3, other: 5, part: 2},
-    {_id: 9, linear: 7, other: 15, part: 1},
-    {_id: 10, linear: null, other: null, part: 2}
+    {_id: 1, linear: 1, other: 1, part: 1, nested: {part: 1}},
+    {_id: 2, linear: 1, other: 1, part: 2, nested: {part: 2}},
+    {_id: 3, linear: null, other: null, part: 1, nested: {part: 1}},
+    {_id: 4, linear: null, other: null, part: 2, nested: {part: 2}},
+    {_id: 5, linear: 5, other: 10, part: 1, nested: {part: 1}},
+    {_id: 6, linear: 6, other: 2, part: 2, nested: {part: 2}},
+    {_id: 7, linear: null, other: null, part: 1, nested: {part: 1}},
+    {_id: 8, linear: 3, other: 5, part: 2, nested: {part: 2}},
+    {_id: 9, linear: 7, other: 15, part: 1, nested: {part: 1}},
+    {_id: 10, linear: null, other: null, part: 2, nested: {part: 2}}
 ];
 
 assert.commandWorked(coll.insert(documents));
@@ -34,7 +34,10 @@ assert.commandWorked(coll.insert(documents));
 const testCases = [
     [
         // Verify $fill partitionBy works with type string and object.
-        [{$fill: {sortBy: {_id: 1}, partitionBy: "$part", output: {linear: {method: "linear"}}}}],
+        [
+            {$project: {nested: 0}},
+            {$fill: {sortBy: {_id: 1}, partitionBy: "$part", output: {linear: {method: "linear"}}}}
+        ],
 
         [
             {_id: 1, linear: 1, other: 1, part: 1},
@@ -52,7 +55,7 @@ const testCases = [
     [
         [
             {$match: {part: 1}},
-            {$project: {other: 0, part: 0}},
+            {$project: {other: 0, part: 0, nested: 0}},
             {$fill: {sortBy: {_id: 1}, output: {linear: {method: "linear"}}}}
         ],
         [
@@ -65,7 +68,7 @@ const testCases = [
     ],  // 1
     [
         [
-            {$project: {linear: 0, part: 0}},
+            {$project: {linear: 0, part: 0, nested: 0}},
             {$fill: {sortBy: {_id: 1}, output: {other: {method: "locf"}}}}
         ],
         [
@@ -83,6 +86,7 @@ const testCases = [
     ],  // 2
     [
         [
+            {$project: {nested: 0}},
             {$match: {part: 2}},
             {
                 $fill: {
@@ -101,10 +105,16 @@ const testCases = [
 
     ],  // 3
     [
-        [{
-            $fill:
-                {sortBy: {_id: 1}, output: {other: {method: "locf"}}, partitionByFields: ["part"]}
-        }],
+        [
+            {$project: {nested: 0}},
+            {
+                $fill: {
+                    sortBy: {_id: 1},
+                    output: {other: {method: "locf"}},
+                    partitionByFields: ["part"]
+                }
+            }
+        ],
         [
             {_id: 1, linear: 1, other: 1, part: 1},
             {_id: 2, linear: 1, other: 1, part: 2},
@@ -120,10 +130,16 @@ const testCases = [
 
     ],  // 4
     [
-        [{
-            $fill:
-                {sortBy: {_id: 1}, output: {other: {method: "locf"}}, partitionBy: {part: "$part"}}
-        }],
+        [
+            {$project: {nested: 0}},
+            {
+                $fill: {
+                    sortBy: {_id: 1},
+                    output: {other: {method: "locf"}},
+                    partitionBy: {part: "$part"}
+                }
+            }
+        ],
         [
             {_id: 1, linear: 1, other: 1, part: 1},
             {_id: 2, linear: 1, other: 1, part: 2},
@@ -139,13 +155,16 @@ const testCases = [
 
     ],  // 5
     [
-        [{
-            $fill: {
-                sortBy: {_id: 1},
-                output: {linear: {method: "linear"}},
-                partitionByFields: ["part"]
+        [
+            {$project: {nested: 0}},
+            {
+                $fill: {
+                    sortBy: {_id: 1},
+                    output: {linear: {method: "linear"}},
+                    partitionByFields: ["part"]
+                }
             }
-        }],
+        ],
         [
             {_id: 1, linear: 1, other: 1, part: 1},
             {_id: 3, linear: 3, other: null, part: 1},
@@ -161,13 +180,16 @@ const testCases = [
 
     ],  // 6
     [
-        [{
-            $fill: {
-                sortBy: {_id: 1},
-                output: {linear: {method: "linear"}},
-                partitionBy: {part: "$part"}
+        [
+            {$project: {nested: 0}},
+            {
+                $fill: {
+                    sortBy: {_id: 1},
+                    output: {linear: {method: "linear"}},
+                    partitionBy: {part: "$part"}
+                }
             }
-        }],
+        ],
         [
             {_id: 1, linear: 1, other: 1, part: 1},
             {_id: 3, linear: 3, other: null, part: 1},
@@ -183,13 +205,16 @@ const testCases = [
 
     ],  // 7
     [
-        [{
-            $fill: {
-                sortBy: {_id: 1},
-                output: {other: {method: "locf"}, linear: {method: "linear"}},
-                partitionByFields: ["part"]
+        [
+            {$project: {nested: 0}},
+            {
+                $fill: {
+                    sortBy: {_id: 1},
+                    output: {other: {method: "locf"}, linear: {method: "linear"}},
+                    partitionByFields: ["part"]
+                }
             }
-        }],
+        ],
         [
             {_id: 1, linear: 1, other: 1, part: 1},
             {_id: 3, linear: 3, other: 1, part: 1},
@@ -206,6 +231,7 @@ const testCases = [
     // Test with first element in partition having a null fill field.
     [
         [
+            {$project: {nested: 0}},
             {$set: {linear: {$cond: [{$eq: ["$linear", 1]}, null, "$linear"]}}},
             {
                 $fill: {
@@ -232,7 +258,7 @@ const testCases = [
     [
         [
             {$match: {part: 1}},
-            {$project: {linear: 0, part: 0}},
+            {$project: {linear: 0, part: 0, nested: 0}},
             {$fill: {sortBy: {_id: 1}, output: {other: {value: "$_id"}}}}
         ],
         [
@@ -246,7 +272,7 @@ const testCases = [
     [
         [
             {$match: {part: 1}},
-            {$project: {linear: 0, part: 0}},
+            {$project: {linear: 0, part: 0, nested: 0}},
             {$fill: {sortBy: {_id: 1}, output: {other: {value: -1}}}}
         ],
         [
@@ -259,6 +285,7 @@ const testCases = [
     ],  // 11
     [
         [
+            {$project: {nested: 0}},
             {$match: {part: 1}},
             {
                 $fill: {
@@ -281,7 +308,7 @@ const testCases = [
     [
         [
             {$match: {part: 1}},
-            {$project: {part: 0}},
+            {$project: {part: 0, nested: 0}},
             {$unionWith: {pipeline: [{$documents: [{_id: 2}, {_id: 4}]}]}},
             {$fill: {sortBy: {_id: 1}, output: {other: {value: "$linear"}}}},
             {$project: {linear: 0}},
@@ -297,6 +324,57 @@ const testCases = [
         ]
 
     ],  // 13
+    // Verify that $fill with 'partitionByFields' can partition by a dotted path.
+    [
+        [
+            {$project: {part: 0, other: 0}},
+            {
+                $fill: {
+                    sortBy: {_id: 1},
+                    partitionByFields: ["nested.part"],
+                    output: {linear: {method: "locf"}}
+                }
+            },
+        ],
+        [
+            {"_id": 1, "linear": 1, "nested": {"part": 1}},
+            {"_id": 2, "linear": 1, "nested": {"part": 2}},
+            {"_id": 3, "linear": 1, "nested": {"part": 1}},
+            {"_id": 4, "linear": 1, "nested": {"part": 2}},
+            {"_id": 5, "linear": 5, "nested": {"part": 1}},
+            {"_id": 6, "linear": 6, "nested": {"part": 2}},
+            {"_id": 7, "linear": 5, "nested": {"part": 1}},
+            {"_id": 8, "linear": 3, "nested": {"part": 2}},
+            {"_id": 9, "linear": 7, "nested": {"part": 1}},
+            {"_id": 10, "linear": 3, "nested": {"part": 2}}
+        ],
+    ],  // 14
+    // Repeat for 'partitionBy'.
+    [
+        [
+            {$project: {part: 0, other: 0}},
+            {
+                $fill: {
+                    sortBy: {_id: 1},
+                    partitionBy: {part: "$nested.part"},
+                    output: {linear: {method: "locf"}}
+                }
+            },
+        ],
+        [
+            {"_id": 1, "linear": 1, "nested": {"part": 1}},
+            {"_id": 2, "linear": 1, "nested": {"part": 2}},
+            {"_id": 3, "linear": 1, "nested": {"part": 1}},
+            {"_id": 4, "linear": 1, "nested": {"part": 2}},
+            {"_id": 5, "linear": 5, "nested": {"part": 1}},
+            {"_id": 6, "linear": 6, "nested": {"part": 2}},
+            {"_id": 7, "linear": 5, "nested": {"part": 1}},
+            {"_id": 8, "linear": 3, "nested": {"part": 2}},
+            {"_id": 9, "linear": 7, "nested": {"part": 1}},
+            {"_id": 10, "linear": 3, "nested": {"part": 2}}
+        ],
+
+    ],  // 15
 
 ];
 
diff --git a/src/mongo/db/pipeline/document_source_fill.cpp b/src/mongo/db/pipeline/document_source_fill.cpp
index e4ad71920d8..768dfcae877 100644
--- a/src/mongo/db/pipeline/document_source_fill.cpp
+++ b/src/mongo/db/pipeline/document_source_fill.cpp
@@ -124,11 +124,11 @@ std::list<boost::intrusive_ptr<DocumentSource>> createFromBson(
     }
 
     if (auto&& partitionByFields = spec.getPartitionByFields()) {
-        BSONObjBuilder partitionBySpec;
+        MutableDocument partitionBySpec;
         for (const auto& fieldName : partitionByFields.value()) {
-            partitionBySpec.append(fieldName, "$" + fieldName);
+            partitionBySpec.setNestedField(fieldName, Value("$" + fieldName));
         }
-        setWindowFieldsSpec.append("partitionBy", partitionBySpec.obj());
+        setWindowFieldsSpec.append("partitionBy", partitionBySpec.freeze().toBson());
     }
 
     std::list<boost::intrusive_ptr<DocumentSource>> finalSources;