SERVER-55217 Treat partitionBy constant as equivalent to no partitioning

2 files changed, 28 insertions, 3 deletions

diff --git a/jstests/aggregation/sources/setWindowFields/partition.js b/jstests/aggregation/sources/setWindowFields/partition.js
index 6d19650edfd..d2cb6be05a3 100644
--- a/jstests/aggregation/sources/setWindowFields/partition.js
+++ b/jstests/aggregation/sources/setWindowFields/partition.js
@@ -1,5 +1,12 @@
 /*
  * Test partitioning inside $setWindowFields.
+ *
+ * @tags: [
+ *   # We assume the pipeline is not split into a shardsPart and mergerPart.
+ *   assumes_unsharded_collection,
+ *   # We're testing the explain plan, not the query results, so the facet passthrough would fail.
+ *   do_not_wrap_aggregations_in_facets,
+ * ]
  */
 (function() {
 "use strict";
@@ -29,4 +36,19 @@ assert.commandFailedWithCode(coll.runCommand({
     cursor: {}
 }),
                              ErrorCodes.TypeMismatch);
+
+// Test that a constant expression for 'partitionBy' is equivalent to no partitioning.
+const constantPartitionExprs = [null, "constant", {$add: [1, 2]}];
+constantPartitionExprs.forEach(function(partitionExpr) {
+    const result = coll.explain().aggregate([
+        // prevent stages from being absorbed into the .find() layer
+        {$_internalInhibitOptimization: {}},
+        {$setWindowFields: {partitionBy: partitionExpr, output: {}}},
+    ]);
+    assert.commandWorked(result);
+    assert(Array.isArray(result.stages), result);
+    assert(result.stages[0].$cursor, result);
+    assert(result.stages[1].$_internalInhibitOptimization, result);
+    assert.eq({$_internalSetWindowFields: {output: {}}}, result.stages[2]);
+});
 })();
\ No newline at end of file
diff --git a/src/mongo/db/pipeline/document_source_set_window_fields.cpp b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
index a84db98aef3..acbd2113a85 100644
--- a/src/mongo/db/pipeline/document_source_set_window_fields.cpp
+++ b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
@@ -154,13 +154,16 @@ list<intrusive_ptr<DocumentSource>> document_source_set_window_fields::create(
     optional<intrusive_ptr<Expression>> complexPartitionBy;
     optional<FieldPath> simplePartitionBy;
     optional<intrusive_ptr<Expression>> simplePartitionByExpr;
-    // If there is no partitionBy, both are empty.
+    // If partitionBy is a constant or there is no partitionBy, both are empty.
     // If partitionBy is already a field path, we only fill in simplePartitionBy.
     // If partitionBy is a more complex expression, we will need to generate a $set stage,
     // which will bind the value of the expression to the name in simplePartitionBy.
     if (partitionBy) {
-        auto exprFieldPath = dynamic_cast<ExpressionFieldPath*>(partitionBy->get());
-        if (exprFieldPath && exprFieldPath->isRootFieldPath()) {
+        partitionBy = (*partitionBy)->optimize();
+        if (dynamic_cast<ExpressionConstant*>(partitionBy->get())) {
+            // partitionBy optimizes to a constant expression, equivalent to a single partition.
+        } else if (auto exprFieldPath = dynamic_cast<ExpressionFieldPath*>(partitionBy->get());
+                   exprFieldPath && exprFieldPath->isRootFieldPath()) {
             // ExpressionFieldPath has "CURRENT" as an explicit first component,
             // but for $sort we don't want that.
             simplePartitionBy = exprFieldPath->getFieldPath().tail();