SERVER-54243 Desugar \$shift window function to \$first with a one doc window

8 files changed, 263 insertions, 42 deletions

diff --git a/src/mongo/db/pipeline/window_function/window_bounds.h b/src/mongo/db/pipeline/window_function/window_bounds.h
index ad6a966ecfa..6999f8fcdbc 100644
--- a/src/mongo/db/pipeline/window_function/window_bounds.h
+++ b/src/mongo/db/pipeline/window_function/window_bounds.h
@@ -86,6 +86,10 @@ struct WindowBounds {
         return WindowBounds{DocumentBased{Unbounded{}, Unbounded{}}};
     }
 
+    static WindowBounds documentBounds(int lower, int upper) {
+        return WindowBounds{DocumentBased{lower, upper}};
+    }
+
     /**
      * Checks whether these bounds are unbounded on both ends.
      * This case is special because it means you don't need a sortBy to interpret the bounds:
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 2bbb520dfa7..7db96ef7e98 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_exec.cpp
@@ -34,6 +34,7 @@
 #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"
 #include "mongo/db/pipeline/window_function/window_function_exec_removable_range.h"
+#include "mongo/db/pipeline/window_function/window_function_shift.h"
 
 namespace mongo {
 
@@ -89,23 +90,24 @@ std::unique_ptr<WindowFunctionExec> translateDocumentWindow(
         bounds.lower);
 }
 
-std::unique_ptr<WindowFunctionExec> translateDerivative(
+std::unique_ptr<mongo::WindowFunctionExec> translateDerivative(
+    window_function::ExpressionDerivative* expr,
     PartitionIterator* iter,
-    const window_function::ExpressionDerivative& deriv,
     const boost::optional<SortPattern>& sortBy) {
-    auto expCtx = deriv.expCtx();
-
     tassert(5490703,
             "$derivative requires a 1-field ascending sortBy",
             sortBy && sortBy->size() == 1 && !sortBy->begin()->expression &&
                 sortBy->begin()->isAscending);
-    auto sortExpr = ExpressionFieldPath::createPathFromString(
-        expCtx, sortBy->begin()->fieldPath->fullPath(), expCtx->variablesParseState);
+    auto sortExpr =
+        ExpressionFieldPath::createPathFromString(expr->expCtx(),
+                                                  sortBy->begin()->fieldPath->fullPath(),
+                                                  expr->expCtx()->variablesParseState);
 
     return std::make_unique<WindowFunctionExecDerivative>(
-        iter, deriv.input(), sortExpr, deriv.bounds(), deriv.outputUnit());
+        iter, expr->input(), sortExpr, expr->bounds(), expr->outputUnit());
 }
 
+
 }  // namespace
 
 std::unique_ptr<WindowFunctionExec> WindowFunctionExec::create(
@@ -114,16 +116,18 @@ std::unique_ptr<WindowFunctionExec> WindowFunctionExec::create(
     const WindowFunctionStatement& functionStmt,
     const boost::optional<SortPattern>& sortBy) {
 
-    if (auto deriv =
-            dynamic_cast<window_function::ExpressionDerivative*>(functionStmt.expr.get())) {
-        return translateDerivative(iter, *deriv, sortBy);
-    } else if (auto first =
+    if (auto expr = dynamic_cast<window_function::ExpressionDerivative*>(functionStmt.expr.get())) {
+        return translateDerivative(expr, iter, sortBy);
+    } else if (auto expr =
                    dynamic_cast<window_function::ExpressionFirst*>(functionStmt.expr.get())) {
-        return std::make_unique<WindowFunctionExecFirst>(
-            iter, first->input(), first->bounds(), boost::none);
-    } else if (auto last =
+        return std::make_unique<WindowFunctionExecFirst>(iter, expr->input(), expr->bounds());
+    } else if (auto expr =
                    dynamic_cast<window_function::ExpressionLast*>(functionStmt.expr.get())) {
-        return std::make_unique<WindowFunctionExecLast>(iter, last->input(), last->bounds());
+        return std::make_unique<WindowFunctionExecLast>(iter, expr->input(), expr->bounds());
+    } else if (auto expr =
+                   dynamic_cast<window_function::ExpressionShift*>(functionStmt.expr.get())) {
+        return std::make_unique<WindowFunctionExecFirst>(
+            iter, expr->input(), expr->bounds(), expr->defaultVal());
     }
 
     WindowBounds bounds = functionStmt.expr->bounds();
diff --git a/src/mongo/db/pipeline/window_function/window_function_exec_first_last.h b/src/mongo/db/pipeline/window_function/window_function_exec_first_last.h
index e5ab3b4065c..775e4aae341 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec_first_last.h
+++ b/src/mongo/db/pipeline/window_function/window_function_exec_first_last.h
@@ -47,19 +47,11 @@ protected:
     WindowFunctionExecForEndpoint(PartitionIterator* iter,
                                   boost::intrusive_ptr<Expression> input,
                                   WindowBounds bounds,
-                                  boost::optional<boost::intrusive_ptr<Expression>> defaultValue)
+                                  const boost::optional<Value>& defaultValue = boost::none)
         : WindowFunctionExec(PartitionAccessor(iter, PartitionAccessor::Policy::kEndpoints)),
           _input(std::move(input)),
-          _bounds(std::move(bounds)) {
-        if (!defaultValue) {
-            _default = Value{BSONNULL};
-        } else {
-            boost::intrusive_ptr<Expression> expr = (*defaultValue)->optimize();
-            ExpressionConstant* ec = dynamic_cast<ExpressionConstant*>(expr.get());
-            tassert(ErrorCodes::FailedToParse, "default expression must be a constant.", ec);
-            _default = ec->getValue();
-        }
-    }
+          _bounds(std::move(bounds)),
+          _default(defaultValue.get_value_or(Value(BSONNULL))) {}
 
     Value getFirst() {
         auto endpoints = _iter.getEndpoints(_bounds);
@@ -90,9 +82,8 @@ public:
     WindowFunctionExecFirst(PartitionIterator* iter,
                             boost::intrusive_ptr<Expression> input,
                             WindowBounds bounds,
-                            boost::optional<boost::intrusive_ptr<Expression>> defaultValue)
-        : WindowFunctionExecForEndpoint(
-              iter, std::move(input), std::move(bounds), std::move(defaultValue)) {}
+                            const boost::optional<Value>& defaultValue = boost::none)
+        : WindowFunctionExecForEndpoint(iter, std::move(input), std::move(bounds), defaultValue) {}
 
     Value getNext() {
         return getFirst();
@@ -104,7 +95,7 @@ public:
     WindowFunctionExecLast(PartitionIterator* iter,
                            boost::intrusive_ptr<Expression> input,
                            WindowBounds bounds)
-        : WindowFunctionExecForEndpoint(iter, input, std::move(bounds), boost::none) {}
+        : WindowFunctionExecForEndpoint(iter, std::move(input), std::move(bounds)) {}
 
     Value getNext() {
         return getLast();
diff --git a/src/mongo/db/pipeline/window_function/window_function_exec_first_last_test.cpp b/src/mongo/db/pipeline/window_function/window_function_exec_first_last_test.cpp
index 3f04f9407b8..07e555c07c6 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec_first_last_test.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_exec_first_last_test.cpp
@@ -48,10 +48,11 @@ public:
     std::pair<WindowFunctionExecFirst, WindowFunctionExecLast> createExecs(
         std::deque<DocumentSource::GetNextResult> docs,
         WindowBounds bounds,
-        boost::optional<BSONElement> defaultValue = boost::none,
+        boost::optional<Value> defaultVal = boost::none,
         boost::optional<std::string> keyPath = boost::none) {
         // 'defaultValue' is an internal functionality of $first needed for $shift desugaring.
         using optExp = boost::optional<boost::intrusive_ptr<Expression>>;
+        using optVal = boost::optional<Value>;
         _docSource = DocumentSourceMock::createForTest(std::move(docs), getExpCtx());
         auto expCtx = getExpCtx().get();
         auto vps = expCtx->variablesParseState;
@@ -63,10 +64,8 @@ public:
                                                     boost::none,
                                                     100 * 1024 * 1024 /* default memory limit */);
         auto inputField = ExpressionFieldPath::parse(expCtx, "$val", vps);
-        auto defaultExp = defaultValue
-            ? optExp(ExpressionConstant::parse(expCtx, *defaultValue, vps))
-            : boost::none;
-        return {WindowFunctionExecFirst(_iter.get(), inputField, bounds, defaultExp),
+
+        return {WindowFunctionExecFirst(_iter.get(), inputField, bounds, defaultVal),
                 WindowFunctionExecLast(_iter.get(), inputField, bounds)};
     }
 
@@ -222,9 +221,7 @@ TEST_F(WindowFunctionExecFirstLastTest, DefaultValueForEmptyWindow1) {
 
     // WindowFunctionExecLast does not accept a default value because it is an internal
     // functionality meant for $shift desugaring.
-    auto [fst, _] = createExecs(std::move(docs),
-                                {WindowBounds::DocumentBased{-3, -2}},
-                                BSON("IGNORED_FIELD_NAME" << -99).firstElement());
+    auto [fst, _] = createExecs(std::move(docs), {WindowBounds::DocumentBased{-3, -2}}, Value(-99));
     ASSERT_VALUE_EQ(Value(-99), fst.getNext());
     advanceIterator();
     ASSERT_VALUE_EQ(Value(-99), fst.getNext());
@@ -242,9 +239,7 @@ TEST_F(WindowFunctionExecFirstLastTest, DefaultValueForEmptyWindow2) {
         Document{{"val", 1}},
     };
 
-    BSONObj defaultValue = BSON("IGNORED_FIELD_NAME" << -99);
-    auto [fst, _] = createExecs(
-        std::move(docs), {WindowBounds::DocumentBased{3, 4}}, defaultValue.firstElement());
+    auto [fst, _] = createExecs(std::move(docs), {WindowBounds::DocumentBased{3, 4}}, Value(-99));
     ASSERT_VALUE_EQ(Value(-99), fst.getNext());
     advanceIterator();
     ASSERT_VALUE_EQ(Value(-99), fst.getNext());
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 49be1652cfc..4618c248ab5 100644
--- a/src/mongo/db/pipeline/window_function/window_function_expression.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_expression.cpp
@@ -37,6 +37,10 @@
 #include "mongo/db/pipeline/lite_parsed_document_source.h"
 #include "mongo/db/query/query_feature_flags_gen.h"
 
+#include "mongo/db/pipeline/window_function/partition_iterator.h"
+#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_expression.h"
 
 using boost::intrusive_ptr;
@@ -185,7 +189,6 @@ boost::intrusive_ptr<Expression> ExpressionFirstLast::parse(
     }
 }
 
-
 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 1f1a66a7807..f03be9bdacf 100644
--- a/src/mongo/db/pipeline/window_function/window_function_expression.h
+++ b/src/mongo/db/pipeline/window_function/window_function_expression.h
@@ -40,6 +40,11 @@
 #include "mongo/db/query/query_feature_flags_gen.h"
 #include "mongo/db/query/sort_pattern.h"
 
+namespace mongo {
+class WindowFunctionExec;
+class PartitionIterator;
+}  // namespace mongo
+
 #define REGISTER_WINDOW_FUNCTION(name, parser)                                       \
     MONGO_INITIALIZER_GENERAL(                                                       \
         addToWindowFunctionMap_##name, ("default"), ("windowFunctionExpressionMap")) \
@@ -54,7 +59,10 @@
         ::mongo::window_function::Expression::registerParser(                                      \
             "$" #name, ::mongo::window_function::ExpressionRemovable<accumClass, wfClass>::parse); \
     }
+
+
 namespace mongo::window_function {
+
 /**
  * A window-function expression describes how to compute a single output value in a
  * $setWindowFields stage. For example, in
@@ -159,6 +167,7 @@ public:
         return args.freezeToValue();
     }
 
+
 protected:
     ExpressionContext* _expCtx;
     std::string _accumulatorName;
diff --git a/src/mongo/db/pipeline/window_function/window_function_shift.cpp b/src/mongo/db/pipeline/window_function/window_function_shift.cpp
new file mode 100644
index 00000000000..026ea0d0977
--- /dev/null
+++ b/src/mongo/db/pipeline/window_function/window_function_shift.cpp
@@ -0,0 +1,133 @@
+/**
+ *    Copyright (C) 2021-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "window_function_shift.h"
+#include "partition_iterator.h"
+#include "window_function_exec_first_last.h"
+
+namespace mongo::window_function {
+REGISTER_WINDOW_FUNCTION(shift, ExpressionShift::parse);
+
+boost::intrusive_ptr<Expression> ExpressionShift::parseShiftArgs(BSONObj obj,
+                                                                 const mongo::StringData& accName,
+                                                                 ExpressionContext* expCtx) {
+    // 'obj' is something like '{output: EXPR, by: INT, default: CONSTEXPR}'.
+    // only default is optional.
+    boost::optional<::mongo::Value> defaultVal;
+    boost::intrusive_ptr<::mongo::Expression> output;
+    int offset;
+
+    bool offsetFound = false;
+    for (const auto& arg : obj) {
+        auto argName = arg.fieldNameStringData();
+
+        if (kOutputArg == argName) {
+            output = ::mongo::Expression::parseOperand(expCtx, arg, expCtx->variablesParseState);
+        } else if (kDefaultArg == argName) {
+            auto defaultExp =
+                ::mongo::Expression::parseOperand(expCtx, arg, expCtx->variablesParseState);
+            defaultExp = defaultExp->optimize();
+            ExpressionConstant* ec = dynamic_cast<ExpressionConstant*>(defaultExp.get());
+            uassert(ErrorCodes::FailedToParse,
+                    str::stream() << "'$shift:" << kDefaultArg
+                                  << "' expression must yield a constant value.",
+                    ec);
+            defaultVal = ec->getValue();
+        } else if (kByArg == argName) {
+            auto parsedOffset = arg.parseIntegerElementToInt();
+            uassert(ErrorCodes::FailedToParse,
+                    str::stream() << "'$shift:" << kByArg
+                                  << "' field must be an integer, but found  " << arg,
+                    parsedOffset.isOK());
+            offset = parsedOffset.getValue();
+            offsetFound = true;
+        } else {
+            uasserted(ErrorCodes::FailedToParse,
+                      str::stream() << "Unknown argument in " << accName);
+        }
+    }
+
+    uassert(ErrorCodes::FailedToParse,
+            str::stream() << accName << " requires an '" << kOutputArg << "' expression.",
+            output);
+
+    uassert(ErrorCodes::FailedToParse,
+            str::stream() << accName << " requires '" << kByArg << "' as an integer value.",
+            offsetFound);
+
+    return make_intrusive<ExpressionShift>(
+        expCtx, accName.toString(), std::move(output), std::move(defaultVal), offset);
+}
+
+boost::intrusive_ptr<Expression> ExpressionShift::parse(BSONObj obj,
+                                                        const boost::optional<SortPattern>& sortBy,
+                                                        ExpressionContext* expCtx) {
+    // 'obj' is something like '{$shift: {<args>}}'.
+    boost::optional<StringData> accumulatorName;
+    boost::intrusive_ptr<Expression> shiftExpr;
+
+    for (const auto& arg : obj) {
+        auto argName = arg.fieldNameStringData();
+        if (argName == kWindowArg) {
+            uasserted(ErrorCodes::FailedToParse,
+                      str::stream() << "$shift does not accept a '" << kWindowArg << "' field");
+        } else if (parserMap.find(argName) != parserMap.end()) {
+            uassert(ErrorCodes::FailedToParse,
+                    "Cannot specify multiple functions in window function spec",
+                    !accumulatorName);
+            accumulatorName = argName;
+            uassert(ErrorCodes::FailedToParse,
+                    "Argument to $shift must be an object",
+                    arg.type() == BSONType::Object);
+            shiftExpr = parseShiftArgs(arg.Obj(), argName, expCtx);
+        } else {
+            uasserted(ErrorCodes::FailedToParse,
+                      str::stream() << "Window function found an unknown argument: " << argName);
+        }
+    }
+
+    uassert(ErrorCodes::FailedToParse,
+            str::stream() << "'" << accumulatorName << "' requires a sortBy",
+            sortBy);
+
+    return shiftExpr;
+}
+
+Value ExpressionShift::serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
+    MutableDocument args;
+    args.addField(kByArg, Value(_offset));
+    args.addField(kOutputArg, _input->serialize(static_cast<bool>(explain)));
+    args.addField(kDefaultArg, _defaultVal.get_value_or(mongo::Value(BSONNULL)));
+
+    MutableDocument windowFun;
+    windowFun.addField(_accumulatorName, args.freezeToValue());
+    return windowFun.freezeToValue();
+}
+
+}  // namespace mongo::window_function
diff --git a/src/mongo/db/pipeline/window_function/window_function_shift.h b/src/mongo/db/pipeline/window_function/window_function_shift.h
new file mode 100644
index 00000000000..99a45dd3f84
--- /dev/null
+++ b/src/mongo/db/pipeline/window_function/window_function_shift.h
@@ -0,0 +1,82 @@
+/**
+ *    Copyright (C) 2021-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/pipeline/window_function/window_function_expression.h"
+
+namespace mongo::window_function {
+
+class ExpressionShift : public Expression {
+public:
+    static constexpr StringData kDefaultArg = "default"_sd;
+    static constexpr StringData kOutputArg = "output"_sd;
+    static constexpr StringData kByArg = "by"_sd;
+
+    static boost::intrusive_ptr<Expression> parse(BSONObj obj,
+                                                  const boost::optional<SortPattern>& sortBy,
+                                                  ExpressionContext* expCtx);
+
+
+    ExpressionShift(ExpressionContext* expCtx,
+                    std::string accumulatorName,
+                    boost::intrusive_ptr<::mongo::Expression> output,
+                    boost::optional<mongo::Value> defaultVal,
+                    int offset)
+        : Expression(expCtx,
+                     accumulatorName,
+                     std::move(output),
+                     WindowBounds::documentBounds(offset, offset)),
+          _defaultVal(std::move(defaultVal)),
+          _offset(offset) {}
+
+    boost::optional<mongo::Value> defaultVal() const {
+        return _defaultVal;
+    }
+
+    boost::intrusive_ptr<AccumulatorState> buildAccumulatorOnly() const final {
+        MONGO_UNREACHABLE_TASSERT(5424301);
+    }
+
+    std::unique_ptr<WindowFunctionState> buildRemovable() const final {
+        MONGO_UNREACHABLE_TASSERT(5424302);
+    }
+
+    Value serialize(boost::optional<ExplainOptions::Verbosity> explain) const final;
+
+private:
+    static boost::intrusive_ptr<Expression> parseShiftArgs(BSONObj obj,
+                                                           const mongo::StringData& accName,
+                                                           ExpressionContext* expCtx);
+
+    boost::optional<mongo::Value> _defaultVal;
+    int _offset;
+};
+
+}  // end namespace mongo::window_function