SERVER-53716 Enable rank accumulator window function translation

3 files changed, 194 insertions, 10 deletions

diff --git a/jstests/aggregation/sources/setWindowFields/rank.js b/jstests/aggregation/sources/setWindowFields/rank.js
new file mode 100644
index 00000000000..c6aec660cfb
--- /dev/null
+++ b/jstests/aggregation/sources/setWindowFields/rank.js
@@ -0,0 +1,119 @@
+/**
+ * Test the rank based window functions.
+ */
+(function() {
+"use strict";
+
+load("jstests/aggregation/extras/utils.js");  // documentEq
+const featureEnabled =
+    assert.commandWorked(db.adminCommand({getParameter: 1, featureFlagWindowFunctions: 1}))
+        .featureFlagWindowFunctions.value;
+if (!featureEnabled) {
+    jsTestLog("Skipping test because the window function feature flag is disabled");
+    return;
+}
+const coll = db[jsTestName()];
+for (let i = 0; i < 12; i++) {
+    coll.insert({_id: i, double: Math.floor(i / 2)});
+}
+
+let origDocs = coll.find().sort({_id: 1});
+function verifyResults(results, valueFunction) {
+    for (let i = 0; i < results.length; i++) {
+        const correctDoc = valueFunction(i, Object.assign({}, origDocs[i]));
+        assert(documentEq(correctDoc, results[i]),
+               "Got: " + tojson(results[i]) + "\nExpected: " + tojson(correctDoc) +
+                   "\n at position " + i + "\n");
+    }
+}
+
+function runRankBasedAccumulator(sortByField, exprField) {
+    return coll
+        .aggregate([{
+            $setWindowFields: {sortBy: sortByField, output: {rank: exprField}},
+        }])
+        .toArray();
+}
+
+// Rank based accumulators don't take windows.
+assert.commandFailedWithCode(coll.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [{
+        $setWindowFields: {
+            sortBy: {_id: 1},
+            output: {rank: {$rank: {}, window: []}},
+        }
+    }],
+    cursor: {}
+}),
+                             5371601);
+
+// Rank based accumulators don't take arguments.
+assert.commandFailedWithCode(coll.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [{
+        $setWindowFields: {
+            sortBy: {_id: 1},
+            output: {rank: {$rank: "$_id"}},
+        }
+    }],
+    cursor: {}
+}),
+                             5371603);
+
+// Rank based accumulators must have a sortBy.
+assert.commandFailedWithCode(coll.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [{
+        $setWindowFields: {
+            output: {rank: {$rank: {}}},
+        }
+    }],
+    cursor: {}
+}),
+                             5371602);
+
+// Rank based accumulators don't take windows.
+assert.commandFailedWithCode(coll.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [{
+        $setWindowFields: {
+            sortBy: {_id: 1},
+            output: {rank: {$rank: {}, window: {documents: [-1, 1]}}},
+        }
+    }],
+    cursor: {}
+}),
+                             5371601);
+
+// Check results with no ties
+let result = runRankBasedAccumulator({_id: 1}, {$rank: {}});
+function noTieFunc(num, baseObj) {
+    baseObj.rank = num + 1;
+    return baseObj;
+}
+verifyResults(result, noTieFunc);
+result = runRankBasedAccumulator({_id: 1}, {$denseRank: {}});
+verifyResults(result, noTieFunc);
+result = runRankBasedAccumulator({_id: 1}, {$documentNumber: {}});
+verifyResults(result, noTieFunc);
+
+// Check results with ties
+origDocs = coll.find().sort({double: 1});
+result = runRankBasedAccumulator({double: 1}, {$rank: {}});
+verifyResults(result, function(num, baseObj) {
+    if (num % 2 == 0) {
+        baseObj.rank = num + 1;
+    } else {
+        baseObj.rank = num;
+    }
+    return baseObj;
+});
+result = runRankBasedAccumulator({double: 1}, {$denseRank: {}});
+verifyResults(result, function(num, baseObj) {
+    baseObj.rank = Math.floor(num / 2) + 1;
+    return baseObj;
+});
+result = runRankBasedAccumulator({double: 1}, {$documentNumber: {}});
+verifyResults(result, noTieFunc);
+})();
\ No newline at end of file
diff --git a/src/mongo/db/pipeline/accumulator_rank.cpp b/src/mongo/db/pipeline/accumulator_rank.cpp
index 26bbc425b23..0ba07f810b7 100644
--- a/src/mongo/db/pipeline/accumulator_rank.cpp
+++ b/src/mongo/db/pipeline/accumulator_rank.cpp
@@ -45,15 +45,13 @@ namespace mongo {
 using boost::intrusive_ptr;
 
 // These don't make sense as accumulators, so only register them as window functions.
-// TODO SERVER-53716 Enable rank function parsing.
-// REGISTER_NON_REMOVABLE_WINDOW_FUNCTION(
-//     rank, mongo::window_function::ExpressionFromRankAccumulator<AccumulatorRank>::parse);
-// REGISTER_NON_REMOVABLE_WINDOW_FUNCTION(
-//     denseRank,
-//     mongo::window_function::ExpressionFromRankAccumulator<AccumulatorDenseRank>::parse);
-// REGISTER_NON_REMOVABLE_WINDOW_FUNCTION(
-//     documentNumber,
-//     mongo::window_function::ExpressionFromRankAccumulator<AccumulatorDocumentNumber>::parse);
+REGISTER_NON_REMOVABLE_WINDOW_FUNCTION(
+    rank, mongo::window_function::ExpressionFromRankAccumulator<AccumulatorRank>::parse);
+REGISTER_NON_REMOVABLE_WINDOW_FUNCTION(
+    denseRank, mongo::window_function::ExpressionFromRankAccumulator<AccumulatorDenseRank>::parse);
+REGISTER_NON_REMOVABLE_WINDOW_FUNCTION(
+    documentNumber,
+    mongo::window_function::ExpressionFromRankAccumulator<AccumulatorDocumentNumber>::parse);
 
 const char* AccumulatorRank::getOpName() const {
     return "$rank";
diff --git a/src/mongo/db/pipeline/window_function/window_function_expression.h b/src/mongo/db/pipeline/window_function/window_function_expression.h
index df6ea29bbb7..ead4de18fc1 100644
--- a/src/mongo/db/pipeline/window_function/window_function_expression.h
+++ b/src/mongo/db/pipeline/window_function/window_function_expression.h
@@ -120,7 +120,7 @@ public:
 
     virtual std::unique_ptr<WindowFunctionState> buildRemovable() const = 0;
 
-    Value serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
+    virtual Value serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
         MutableDocument args;
 
         args[_accumulatorName] = _input->serialize(static_cast<bool>(explain));
@@ -244,4 +244,71 @@ public:
     }
 };
 
+template <typename RankType>
+class ExpressionFromRankAccumulator : public Expression {
+public:
+    static boost::intrusive_ptr<Expression> parse(BSONObj obj,
+                                                  const boost::optional<SortPattern>& sortBy,
+                                                  ExpressionContext* expCtx) {
+        // 'obj' is something like '{$func: <args>}'
+        uassert(5371601, "Rank style window functions take no other arguments", obj.nFields() == 1);
+        boost::optional<StringData> accumulatorName;
+        // Rank based accumulators are always unbounded to current.
+        WindowBounds bounds = WindowBounds{
+            WindowBounds::DocumentBased{WindowBounds::Unbounded{}, WindowBounds::Current{}}};
+        boost::intrusive_ptr<::mongo::Expression> input;
+        auto arg = obj.firstElement();
+        auto argName = arg.fieldNameStringData();
+        if (parserMap.find(argName) != parserMap.end()) {
+            uassert(5371603,
+                    str::stream() << accumulatorName << " must be specified with '{}' as the value",
+                    arg.type() == BSONType::Object && arg.embeddedObject().nFields() == 0);
+            accumulatorName = argName;
+        } else {
+            tasserted(ErrorCodes::FailedToParse,
+                      str::stream() << "Window function found an unknown argument: " << argName);
+        }
+
+        // Rank based accumulators use the sort by expression as the input.
+        uassert(
+            5371602,
+            str::stream()
+                << accumulatorName
+                << " must be specified with a top level sortBy expression with exactly one element",
+            sortBy && sortBy->isSingleElementKey());
+        auto sortPatternPart = sortBy.get()[0];
+        if (sortPatternPart.fieldPath) {
+            auto sortExpression = ExpressionFieldPath::createPathFromString(
+                expCtx, sortPatternPart.fieldPath->fullPath(), expCtx->variablesParseState);
+            return make_intrusive<ExpressionFromRankAccumulator<RankType>>(
+                expCtx, accumulatorName->toString(), std::move(sortExpression), std::move(bounds));
+        } else {
+            return make_intrusive<ExpressionFromRankAccumulator<RankType>>(
+                expCtx, accumulatorName->toString(), sortPatternPart.expression, std::move(bounds));
+        }
+    }
+
+    ExpressionFromRankAccumulator(ExpressionContext* expCtx,
+                                  std::string accumulatorName,
+                                  boost::intrusive_ptr<::mongo::Expression> input,
+                                  WindowBounds bounds)
+        : Expression(expCtx, std::move(accumulatorName), std::move(input), std::move(bounds)) {}
+
+    boost::intrusive_ptr<AccumulatorState> buildAccumulatorOnly() const final {
+        return RankType::create(_expCtx);
+    }
+
+    std::unique_ptr<WindowFunctionState> buildRemovable() const final {
+        tasserted(5371600,
+                  str::stream() << "Window function " << _accumulatorName
+                                << " is not supported with a removable window");
+    }
+
+    Value serialize(boost::optional<ExplainOptions::Verbosity> explain) const final {
+        MutableDocument args;
+        args.addField(_accumulatorName, Value(Document()));
+        return args.freezeToValue();
+    }
+};
+
 }  // namespace mongo::window_function