SERVER-57886 add top/bottom/topN/bottomN operators as window functions

1 files changed, 90 insertions, 20 deletions

diff --git a/jstests/aggregation/sources/setWindowFields/n_accumulators.js b/jstests/aggregation/sources/setWindowFields/n_accumulators.js
index 20b5cabd75d..4fb2c1242be 100644
--- a/jstests/aggregation/sources/setWindowFields/n_accumulators.js
+++ b/jstests/aggregation/sources/setWindowFields/n_accumulators.js
@@ -9,17 +9,42 @@ load("jstests/aggregation/extras/window_function_helpers.js");
 const coll = db[jsTestName()];
 coll.drop();
 
-// TODO SERVER-57886: Add test cases for $top/$bottom/$topN/$bottomN window functions.
-const nAccumulators = ["$minN", "$maxN", "$firstN", "$lastN"];
+const needsSortBy = (op) => ({
+    $minN: false,
+    $maxN: false,
+    $firstN: false,
+    $lastN: false,
+    $topN: true,
+    $bottomN: true,
+    $top: true,
+    $bottom: true,
+}[op]);
+
+// A map from the accumulator name to a function that ignores the parameters it doesn't need and
+// generates a valid accumulator spec.
+const simple = (n, input) => ({n, input: "$" + input});
+const topBottomN = (n, output) => ({n, output: "$" + output, sortBy: {[output]: 1}});
+const topBottom = (n, output) => ({output: "$" + output, sortBy: {[output]: 1}});
+const nAccumulators = {
+    $minN: simple,
+    $maxN: simple,
+    $firstN: simple,
+    $lastN: simple,
+    $topN: topBottomN,
+    $bottomN: topBottomN,
+    $top: topBottom,
+    $bottom: topBottom,
+};
 
 // Create a collection of tickers and prices.
 const nDocsPerTicker = 10;
 seedWithTickerData(coll, nDocsPerTicker);
 
-for (const acc of nAccumulators) {
+for (const acc of Object.keys(nAccumulators)) {
     for (const nValue of [4, 7, 12]) {
         jsTestLog("Testing accumulator " + tojson(acc) + " with 'n' set to " + tojson(nValue));
-        testAccumAgainstGroup(coll, acc, [], {input: "$price", n: nValue});
+        const noValue = (acc === "$top" || acc === "$bottom") ? null : [];
+        testAccumAgainstGroup(coll, acc, noValue, nAccumulators[acc](nValue, "price"));
     }
 
     // Verify that the accumulator will not throw if the 'n' expression evaluates to a constant.
@@ -28,7 +53,7 @@ for (const acc of nAccumulators) {
             $setWindowFields: {
                 partitionBy: "$ticker",
                 sortBy: {_id: 1},
-                output: {res: {[acc]: {n: {$add: [1, 2]}, input: "$price"}}}
+                output: {res: {[acc]: nAccumulators[acc]({$add: [1, 2]}, "price")}}
             },
         },
     ];
@@ -39,6 +64,7 @@ for (const acc of nAccumulators) {
     function testError(accSpec, expectedCode) {
         assert.throwsWithCode(() => coll.aggregate([{
             $setWindowFields: {
+                partitionBy: "$ticket",
                 sortBy: {ts: 1},
                 output: {outputField: accSpec},
             }
@@ -50,23 +76,67 @@ for (const acc of nAccumulators) {
     // parsing error due to a non-object specification.
     const expectedCode = (acc === "$minN" || acc === "$maxN") ? 5787900 : 5787801;
 
-    // Invalid/missing accumulator specification.
-    testError({[acc]: "non object"}, expectedCode);
-    testError({window: {documents: [-1, 1]}}, ErrorCodes.FailedToParse);
-    testError({[acc]: {n: 2}, window: {documents: [-1, 1]}}, 5787907);
-    testError({[acc]: {input: "$foo"}, window: {documents: [-1, 1]}}, 5787906);
-    testError({[acc]: {input: "$foo", n: 2.1}, window: {documents: [-1, 1]}}, 5787903);
+    if (!needsSortBy(acc)) {
+        // Invalid/missing accumulator specification.
+        testError({[acc]: "non object"}, expectedCode);
+        testError({window: {documents: [-1, 1]}}, ErrorCodes.FailedToParse);
+        testError({[acc]: {n: 2}, window: {documents: [-1, 1]}}, 5787907);
+        testError({[acc]: {input: "$foo"}, window: {documents: [-1, 1]}}, 5787906);
+        testError({[acc]: {input: "$foo", n: 2.1}, window: {documents: [-1, 1]}}, 5787903);
+
+        // Invalid window specification.
+        testError({[acc]: {input: "$foo", n: 2.0}, window: [-1, 1]}, ErrorCodes.FailedToParse);
+
+        // Non constant argument for 'n'.
+        testError({[acc]: {input: "$foo", n: "$a"}, window: {documents: [-1, 1]}}, 5787902);
+
+        // Can't reference partition key.
+        testError({[acc]: {input: "$foo", n: "$ticket"}, window: {documents: [-1, 1]}}, 5787902);
+
+        // n = 0
+        testError({[acc]: {input: "$foo", n: 0}, window: {documents: [-1, 1]}}, 5787908);
 
-    // Invalid window specification.
-    testError({[acc]: {input: "$foo", n: 2.0}, window: [-1, 1]}, ErrorCodes.FailedToParse);
+        // n < 0
+        testError({[acc]: {input: "$foo", n: -100}, window: {documents: [-1, 1]}}, 5787908);
+    } else if (acc == "$topN" || acc == "$bottomN") {
+        // TODO SERVER-59327 combine error codes if we decide to use the same parser function.
+        // Invalid/missing accumulator specification.
+        const sortBy = {"foo": 1};
+        const output = "$foo";
+        testError({[acc]: "non object"}, 5788001);
+        testError({window: {documents: [-1, 1]}}, ErrorCodes.FailedToParse);
+        testError({[acc]: {n: 2, sortBy}, window: {documents: [-1, 1]}}, 5788004);
+        testError({[acc]: {output, sortBy}, window: {documents: [-1, 1]}}, 5788003);
+        testError({[acc]: {output, sortBy, n: 2.1}, window: {documents: [-1, 1]}}, 5787903);
 
-    // Non constant argument for 'n'.
-    testError({[acc]: {input: "$foo", n: "$a"}, window: {documents: [-1, 1]}}, 5787902);
+        // Missing sortBy.
+        testError({[acc]: {output, n: 2}, window: {documents: [-1, 1]}}, 5788005);
 
-    // n = 0
-    testError({[acc]: {input: "$foo", n: 0}, window: {documents: [-1, 1]}}, 5787908);
+        // Invalid window specification.
+        testError({[acc]: {output, sortBy, n: 2.0}, window: [-1, 1]}, ErrorCodes.FailedToParse);
 
-    // n < 0
-    testError({[acc]: {input: "$foo", n: -100}, window: {documents: [-1, 1]}}, 5787908);
+        // Non constant argument for 'n'.
+        testError({[acc]: {output, sortBy, n: "$a"}, window: {documents: [-1, 1]}}, 5787902);
+
+        // Can't reference partition key.
+        testError({[acc]: {output, sortBy, n: "$ticket"}, window: {documents: [-1, 1]}}, 5787902);
+
+        // n = 0
+        testError({[acc]: {output, sortBy, n: 0}, window: {documents: [-1, 1]}}, 5787908);
+
+        // n < 0
+        testError({[acc]: {output, sortBy, n: -100}, window: {documents: [-1, 1]}}, 5787908);
+    } else {
+        // $top/$bottom parsing tests.
+        const sortBy = {"foo": 1};
+        const output = "$foo";
+        testError({[acc]: "non object"}, 5788001);
+        testError({window: {documents: [-1, 1]}}, ErrorCodes.FailedToParse);
+        testError({[acc]: {sortBy}, window: {documents: [-1, 1]}}, 5788004);
+        testError({[acc]: {n: 2, output, sortBy}, window: {documents: [-1, 1]}}, 5788002);
+
+        // Missing sortBy.
+        testError({[acc]: {output}, window: {documents: [-1, 1]}}, 5788005);
+    }
 }
-})();
\ No newline at end of file
+})();