SERVER-55064 Add translation logic for $first/$last executor

5 files changed, 339 insertions, 0 deletions

diff --git a/jstests/aggregation/sources/setWindowFields/first.js b/jstests/aggregation/sources/setWindowFields/first.js
new file mode 100644
index 00000000000..e5043250d7e
--- /dev/null
+++ b/jstests/aggregation/sources/setWindowFields/first.js
@@ -0,0 +1,102 @@
+/**
+ * Test the behavior of $first.
+ */
+(function() {
+"use strict";
+
+load("jstests/aggregation/extras/window_function_helpers.js");
+
+const getParam = db.adminCommand({getParameter: 1, featureFlagWindowFunctions: 1});
+jsTestLog(getParam);
+const featureEnabled = assert.commandWorked(getParam).featureFlagWindowFunctions.value;
+if (!featureEnabled) {
+    jsTestLog("Skipping test because the window function feature flag is disabled");
+    return;
+}
+
+const coll = db[jsTestName()];
+coll.drop();
+
+// Create a collection of tickers and prices.
+const nDocsPerTicker = 10;
+seedWithTickerData(coll, nDocsPerTicker);
+
+// Run the suite of partition and bounds tests against the $first function.
+testAccumAgainstGroup(coll, "$first");
+
+// Like most other window functions, the default window for $first is [unbounded, unbounded].
+coll.drop();
+assert.commandWorked(coll.insert([
+    {x: 0, y: 0},
+    {x: 1, y: 42},
+    {x: 2, y: 67},
+    {x: 3, y: 99},
+    {x: 4, y: 20},
+]));
+let result = coll.aggregate([
+                     {
+                         $setWindowFields: {
+                             sortBy: {x: 1},
+                             output: {
+                                 first: {$first: "$y"},
+                             }
+                         }
+                     },
+                     {$unset: "_id"},
+                 ])
+                 .toArray();
+assert.sameMembers(result, [
+    {x: 0, y: 0, first: 0},
+    {x: 1, y: 42, first: 0},
+    {x: 2, y: 67, first: 0},
+    {x: 3, y: 99, first: 0},
+    {x: 4, y: 20, first: 0},
+]);
+
+// A default value of NULL is returned if there is no first document.
+coll.drop();
+assert.commandWorked(coll.insert([
+    {x: 1, y: 5},
+    {x: 2, y: 4},
+    {x: 3, y: 6},
+    {x: 4, y: 5},
+]));
+result = coll.aggregate([
+                 {
+                     $setWindowFields: {
+                         sortBy: {x: 1},
+                         partitionBy: "$x",
+                         output: {
+                             first: {$first: "$y", window: {documents: [-1, -1]}},
+                         }
+                     }
+                 },
+                 {$unset: "_id"},
+             ])
+             .toArray();
+assert.sameMembers(result, [
+    {x: 1, y: 5, first: null},
+    {x: 2, y: 4, first: null},
+    {x: 3, y: 6, first: null},
+    {x: 4, y: 5, first: null},
+]);
+
+// Nonobject window fields cause parse errors
+result = coll.runCommand({
+    explain: {
+        aggregate: coll.getName(),
+        cursor: {},
+        pipeline: [
+            {
+                $setWindowFields: {
+                    sortBy: {x: 1},
+                    output: {
+                        first: {$first: "$y", window: [0, 1]},
+                    }
+                }
+            },
+        ]
+    }
+});
+assert.commandFailedWithCode(result, ErrorCodes.FailedToParse, "'window' field must be an object");
+})();
diff --git a/jstests/aggregation/sources/setWindowFields/last.js b/jstests/aggregation/sources/setWindowFields/last.js
new file mode 100644
index 00000000000..b974497fb40
--- /dev/null
+++ b/jstests/aggregation/sources/setWindowFields/last.js
@@ -0,0 +1,102 @@
+/**
+ * Test the behavior of $last.
+ */
+(function() {
+"use strict";
+
+load("jstests/aggregation/extras/window_function_helpers.js");
+
+const getParam = db.adminCommand({getParameter: 1, featureFlagWindowFunctions: 1});
+jsTestLog(getParam);
+const featureEnabled = assert.commandWorked(getParam).featureFlagWindowFunctions.value;
+if (!featureEnabled) {
+    jsTestLog("Skipping test because the window function feature flag is disabled");
+    return;
+}
+
+const coll = db[jsTestName()];
+coll.drop();
+
+// Create a collection of tickers and prices.
+const nDocsPerTicker = 10;
+seedWithTickerData(coll, nDocsPerTicker);
+
+// Run the suite of partition and bounds tests against the $last function.
+testAccumAgainstGroup(coll, "$last");
+
+// Like most other window functions, the default window for $last is [unbounded, unbounded].
+coll.drop();
+assert.commandWorked(coll.insert([
+    {x: 0, y: 0},
+    {x: 1, y: 42},
+    {x: 2, y: 67},
+    {x: 3, y: 99},
+    {x: 4, y: 20},
+]));
+let result = coll.aggregate([
+                     {
+                         $setWindowFields: {
+                             sortBy: {x: 1},
+                             output: {
+                                 last: {$last: "$y"},
+                             }
+                         }
+                     },
+                     {$unset: "_id"},
+                 ])
+                 .toArray();
+assert.sameMembers(result, [
+    {x: 0, y: 0, last: 20},
+    {x: 1, y: 42, last: 20},
+    {x: 2, y: 67, last: 20},
+    {x: 3, y: 99, last: 20},
+    {x: 4, y: 20, last: 20},
+]);
+
+// A default value of NULL is returned if there is no last document.
+coll.drop();
+assert.commandWorked(coll.insert([
+    {x: 1, y: 5},
+    {x: 2, y: 4},
+    {x: 3, y: 6},
+    {x: 4, y: 5},
+]));
+result = coll.aggregate([
+                 {
+                     $setWindowFields: {
+                         sortBy: {x: 1},
+                         partitionBy: "$x",
+                         output: {
+                             last: {$last: "$y", window: {documents: [-1, -1]}},
+                         }
+                     }
+                 },
+                 {$unset: "_id"},
+             ])
+             .toArray();
+assert.sameMembers(result, [
+    {x: 1, y: 5, last: null},
+    {x: 2, y: 4, last: null},
+    {x: 3, y: 6, last: null},
+    {x: 4, y: 5, last: null},
+]);
+
+// Nonobject window fields cause parse errors
+result = coll.runCommand({
+    explain: {
+        aggregate: coll.getName(),
+        cursor: {},
+        pipeline: [
+            {
+                $setWindowFields: {
+                    sortBy: {x: 1},
+                    output: {
+                        last: {$last: "$y", window: [0, 1]},
+                    }
+                }
+            },
+        ]
+    }
+});
+assert.commandFailedWithCode(result, ErrorCodes.FailedToParse, "'window' field must be an object");
+})();
diff --git a/src/mongo/db/pipeline/window_function/window_function_exec.cpp b/src/mongo/db/pipeline/window_function/window_function_exec.cpp
index a3df28b50ea..0f144f930da 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_exec.cpp
@@ -29,6 +29,7 @@
 
 #include "mongo/db/pipeline/window_function/window_function_exec.h"
 #include "mongo/db/pipeline/window_function/window_function_exec_derivative.h"
+#include "mongo/db/pipeline/window_function/window_function_exec_first_last.h"
 #include "mongo/db/pipeline/window_function/window_function_exec_non_removable.h"
 #include "mongo/db/pipeline/window_function/window_function_exec_non_removable_range.h"
 #include "mongo/db/pipeline/window_function/window_function_exec_removable_document.h"
@@ -116,6 +117,13 @@ std::unique_ptr<WindowFunctionExec> WindowFunctionExec::create(
     if (auto deriv =
             dynamic_cast<window_function::ExpressionDerivative*>(functionStmt.expr.get())) {
         return translateDerivative(iter, *deriv, sortBy);
+    } else if (auto first =
+                   dynamic_cast<window_function::ExpressionFirst*>(functionStmt.expr.get())) {
+        return std::make_unique<WindowFunctionExecFirst>(
+            iter, first->input(), first->bounds(), boost::none);
+    } else if (auto last =
+                   dynamic_cast<window_function::ExpressionLast*>(functionStmt.expr.get())) {
+        return std::make_unique<WindowFunctionExecLast>(iter, last->input(), last->bounds());
     }
 
     WindowBounds bounds = functionStmt.expr->bounds();
diff --git a/src/mongo/db/pipeline/window_function/window_function_expression.cpp b/src/mongo/db/pipeline/window_function/window_function_expression.cpp
index e9902f42366..df551b0db93 100644
--- a/src/mongo/db/pipeline/window_function/window_function_expression.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_expression.cpp
@@ -45,6 +45,8 @@ using boost::optional;
 namespace mongo::window_function {
 
 REGISTER_WINDOW_FUNCTION(derivative, ExpressionDerivative::parse);
+REGISTER_WINDOW_FUNCTION(first, ExpressionFirst::parse);
+REGISTER_WINDOW_FUNCTION(last, ExpressionLast::parse);
 
 StringMap<Expression::Parser> Expression::parserMap;
 
@@ -124,6 +126,63 @@ boost::intrusive_ptr<Expression> ExpressionExpMovingAvg::parse(
     }
 }
 
+boost::intrusive_ptr<Expression> ExpressionFirstLast::parse(
+    BSONObj obj,
+    const boost::optional<SortPattern>& sortBy,
+    ExpressionContext* expCtx,
+    Sense sense) {
+    // Example document:
+    // {
+    //   accumulatorName: <expr>,
+    //   window: {...} // optional
+    // }
+
+    const std::string& accumulatorName = senseToAccumulatorName(sense);
+    boost::optional<WindowBounds> bounds;
+    boost::intrusive_ptr<::mongo::Expression> input;
+    for (const auto& arg : obj) {
+        auto argName = arg.fieldNameStringData();
+        if (argName == kWindowArg) {
+            uassert(ErrorCodes::FailedToParse,
+                    "'window' field must be an object",
+                    obj[kWindowArg].type() == BSONType::Object);
+            uassert(ErrorCodes::FailedToParse,
+                    str::stream() << "saw multiple 'window' fields in '" << accumulatorName
+                                  << "' expression",
+                    bounds == boost::none);
+            bounds = WindowBounds::parse(arg.embeddedObject(), sortBy, expCtx);
+        } else if (argName == StringData(accumulatorName)) {
+            input = ::mongo::Expression::parseOperand(expCtx, arg, expCtx->variablesParseState);
+
+        } else {
+            uasserted(ErrorCodes::FailedToParse,
+                      str::stream() << accumulatorName << " got unexpected argument: " << argName);
+        }
+    }
+    tassert(ErrorCodes::FailedToParse,
+            str::stream() << accumulatorName << " parser called with no " << accumulatorName
+                          << " key",
+            input);
+
+    // The default window bounds are [unbounded, unbounded].
+    if (!bounds) {
+        bounds = WindowBounds{
+            WindowBounds::DocumentBased{WindowBounds::Unbounded{}, WindowBounds::Unbounded{}}};
+    }
+
+    switch (sense) {
+        case Sense::kFirst:
+            return make_intrusive<ExpressionFirst>(expCtx, std::move(input), std::move(*bounds));
+        case Sense::kLast:
+            return make_intrusive<ExpressionLast>(expCtx, std::move(input), std::move(*bounds));
+        default:
+            uasserted(ErrorCodes::FailedToParse,
+                      str::stream() << accumulatorName << " is not $first or $last");
+            return nullptr;
+    }
+}
+
+
 MONGO_INITIALIZER(windowFunctionExpressionMap)(InitializerContext*) {
     // Nothing to do. This initializer exists to tie together all the individual initializers
     // defined by REGISTER_WINDOW_FUNCTION and REGISTER_REMOVABLE_WINDOW_FUNCTION
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 a34bc68c0d5..c74d274f142 100644
--- a/src/mongo/db/pipeline/window_function/window_function_expression.h
+++ b/src/mongo/db/pipeline/window_function/window_function_expression.h
@@ -643,4 +643,72 @@ public:
     }
 };
 
+
+class ExpressionFirstLast : public Expression {
+public:
+    enum Sense : int {
+        kFirst,
+        kLast,
+    };
+
+    static boost::intrusive_ptr<Expression> parse(BSONObj obj,
+                                                  const boost::optional<SortPattern>& sortBy,
+                                                  ExpressionContext* expCtx,
+                                                  Sense sense);
+    static std::string senseToAccumulatorName(Sense sense) {
+        switch (sense) {
+            case Sense::kFirst:
+                return "$first";
+            case Sense::kLast:
+                return "$last";
+            default:
+                return "unrecognized sense";
+        }
+    }
+};
+
+class ExpressionFirst : public Expression {
+public:
+    ExpressionFirst(ExpressionContext* expCtx,
+                    boost::intrusive_ptr<::mongo::Expression> input,
+                    WindowBounds bounds)
+        : Expression(expCtx, "$first", std::move(input), std::move(bounds)) {}
+
+    static boost::intrusive_ptr<Expression> parse(BSONObj obj,
+                                                  const boost::optional<SortPattern>& sortBy,
+                                                  ExpressionContext* expCtx) {
+        return ExpressionFirstLast::parse(obj, sortBy, expCtx, ExpressionFirstLast::Sense::kFirst);
+    }
+
+    boost::intrusive_ptr<AccumulatorState> buildAccumulatorOnly() const final {
+        MONGO_UNREACHABLE_TASSERT(5490701);
+    }
+
+    std::unique_ptr<WindowFunctionState> buildRemovable() const final {
+        MONGO_UNREACHABLE_TASSERT(5490702);
+    }
+};
+
+class ExpressionLast : public Expression {
+public:
+    ExpressionLast(ExpressionContext* expCtx,
+                   boost::intrusive_ptr<::mongo::Expression> input,
+                   WindowBounds bounds)
+        : Expression(expCtx, "$last", std::move(input), std::move(bounds)) {}
+
+    static boost::intrusive_ptr<Expression> parse(BSONObj obj,
+                                                  const boost::optional<SortPattern>& sortBy,
+                                                  ExpressionContext* expCtx) {
+        return ExpressionFirstLast::parse(obj, sortBy, expCtx, ExpressionFirstLast::Sense::kLast);
+    }
+
+    boost::intrusive_ptr<AccumulatorState> buildAccumulatorOnly() const final {
+        MONGO_UNREACHABLE_TASSERT(5490701);
+    }
+
+    std::unique_ptr<WindowFunctionState> buildRemovable() const final {
+        MONGO_UNREACHABLE_TASSERT(5490702);
+    }
+};
+
 }  // namespace mongo::window_function