From ffb78b55cfeea62fd76fada89c28fc2403b6c2d5 Mon Sep 17 00:00:00 2001 From: Joel Redman Date: Wed, 21 Apr 2021 22:24:25 +0000 Subject: SERVER-54243 Desugar \$shift window function to \$first with a one doc window --- .../db/pipeline/window_function/window_bounds.h | 4 + .../window_function/window_function_exec.cpp | 34 +++--- .../window_function_exec_first_last.h | 21 +--- .../window_function_exec_first_last_test.cpp | 17 +-- .../window_function/window_function_expression.cpp | 5 +- .../window_function/window_function_expression.h | 9 ++ .../window_function/window_function_shift.cpp | 133 +++++++++++++++++++++ .../window_function/window_function_shift.h | 82 +++++++++++++ 8 files changed, 263 insertions(+), 42 deletions(-) create mode 100644 src/mongo/db/pipeline/window_function/window_function_shift.cpp create mode 100644 src/mongo/db/pipeline/window_function/window_function_shift.h (limited to 'src/mongo/db/pipeline/window_function') 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 translateDocumentWindow( bounds.lower); } -std::unique_ptr translateDerivative( +std::unique_ptr translateDerivative( + window_function::ExpressionDerivative* expr, PartitionIterator* iter, - const window_function::ExpressionDerivative& deriv, const boost::optional& 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( - iter, deriv.input(), sortExpr, deriv.bounds(), deriv.outputUnit()); + iter, expr->input(), sortExpr, expr->bounds(), expr->outputUnit()); } + } // namespace std::unique_ptr WindowFunctionExec::create( @@ -114,16 +116,18 @@ std::unique_ptr WindowFunctionExec::create( const WindowFunctionStatement& functionStmt, const boost::optional& sortBy) { - if (auto deriv = - dynamic_cast(functionStmt.expr.get())) { - return translateDerivative(iter, *deriv, sortBy); - } else if (auto first = + if (auto expr = dynamic_cast(functionStmt.expr.get())) { + return translateDerivative(expr, iter, sortBy); + } else if (auto expr = dynamic_cast(functionStmt.expr.get())) { - return std::make_unique( - iter, first->input(), first->bounds(), boost::none); - } else if (auto last = + return std::make_unique(iter, expr->input(), expr->bounds()); + } else if (auto expr = dynamic_cast(functionStmt.expr.get())) { - return std::make_unique(iter, last->input(), last->bounds()); + return std::make_unique(iter, expr->input(), expr->bounds()); + } else if (auto expr = + dynamic_cast(functionStmt.expr.get())) { + return std::make_unique( + 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 input, WindowBounds bounds, - boost::optional> defaultValue) + const boost::optional& 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 expr = (*defaultValue)->optimize(); - ExpressionConstant* ec = dynamic_cast(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 input, WindowBounds bounds, - boost::optional> defaultValue) - : WindowFunctionExecForEndpoint( - iter, std::move(input), std::move(bounds), std::move(defaultValue)) {} + const boost::optional& 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 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 createExecs( std::deque docs, WindowBounds bounds, - boost::optional defaultValue = boost::none, + boost::optional defaultVal = boost::none, boost::optional keyPath = boost::none) { // 'defaultValue' is an internal functionality of $first needed for $shift desugaring. using optExp = boost::optional>; + using optVal = boost::optional; _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 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::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 + * . + * + * 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 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(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( + expCtx, accName.toString(), std::move(output), std::move(defaultVal), offset); +} + +boost::intrusive_ptr ExpressionShift::parse(BSONObj obj, + const boost::optional& sortBy, + ExpressionContext* expCtx) { + // 'obj' is something like '{$shift: {}}'. + boost::optional accumulatorName; + boost::intrusive_ptr 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 explain) const { + MutableDocument args; + args.addField(kByArg, Value(_offset)); + args.addField(kOutputArg, _input->serialize(static_cast(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 + * . + * + * 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 parse(BSONObj obj, + const boost::optional& sortBy, + ExpressionContext* expCtx); + + + ExpressionShift(ExpressionContext* expCtx, + std::string accumulatorName, + boost::intrusive_ptr<::mongo::Expression> output, + boost::optional defaultVal, + int offset) + : Expression(expCtx, + accumulatorName, + std::move(output), + WindowBounds::documentBounds(offset, offset)), + _defaultVal(std::move(defaultVal)), + _offset(offset) {} + + boost::optional defaultVal() const { + return _defaultVal; + } + + boost::intrusive_ptr buildAccumulatorOnly() const final { + MONGO_UNREACHABLE_TASSERT(5424301); + } + + std::unique_ptr buildRemovable() const final { + MONGO_UNREACHABLE_TASSERT(5424302); + } + + Value serialize(boost::optional explain) const final; + +private: + static boost::intrusive_ptr parseShiftArgs(BSONObj obj, + const mongo::StringData& accName, + ExpressionContext* expCtx); + + boost::optional _defaultVal; + int _offset; +}; + +} // end namespace mongo::window_function -- cgit v1.2.1