From 9fde6f9838a46a0ad98e6391bdc1e8d2c845c448 Mon Sep 17 00:00:00 2001 From: samontea Date: Wed, 14 Apr 2021 16:15:35 +0000 Subject: SERVER-55064 Add translation logic for $first/$last executor --- .../aggregation/sources/setWindowFields/first.js | 102 +++++++++++++++++++++ .../aggregation/sources/setWindowFields/last.js | 102 +++++++++++++++++++++ .../window_function/window_function_exec.cpp | 8 ++ .../window_function/window_function_expression.cpp | 59 ++++++++++++ .../window_function/window_function_expression.h | 68 ++++++++++++++ 5 files changed, 339 insertions(+) create mode 100644 jstests/aggregation/sources/setWindowFields/first.js create mode 100644 jstests/aggregation/sources/setWindowFields/last.js 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::create( if (auto deriv = dynamic_cast(functionStmt.expr.get())) { return translateDerivative(iter, *deriv, sortBy); + } else if (auto first = + dynamic_cast(functionStmt.expr.get())) { + return std::make_unique( + iter, first->input(), first->bounds(), boost::none); + } else if (auto last = + dynamic_cast(functionStmt.expr.get())) { + return std::make_unique(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::parserMap; @@ -124,6 +126,63 @@ boost::intrusive_ptr ExpressionExpMovingAvg::parse( } } +boost::intrusive_ptr ExpressionFirstLast::parse( + BSONObj obj, + const boost::optional& sortBy, + ExpressionContext* expCtx, + Sense sense) { + // Example document: + // { + // accumulatorName: , + // window: {...} // optional + // } + + const std::string& accumulatorName = senseToAccumulatorName(sense); + boost::optional 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(expCtx, std::move(input), std::move(*bounds)); + case Sense::kLast: + return make_intrusive(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 parse(BSONObj obj, + const boost::optional& 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 parse(BSONObj obj, + const boost::optional& sortBy, + ExpressionContext* expCtx) { + return ExpressionFirstLast::parse(obj, sortBy, expCtx, ExpressionFirstLast::Sense::kFirst); + } + + boost::intrusive_ptr buildAccumulatorOnly() const final { + MONGO_UNREACHABLE_TASSERT(5490701); + } + + std::unique_ptr 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 parse(BSONObj obj, + const boost::optional& sortBy, + ExpressionContext* expCtx) { + return ExpressionFirstLast::parse(obj, sortBy, expCtx, ExpressionFirstLast::Sense::kLast); + } + + boost::intrusive_ptr buildAccumulatorOnly() const final { + MONGO_UNREACHABLE_TASSERT(5490701); + } + + std::unique_ptr buildRemovable() const final { + MONGO_UNREACHABLE_TASSERT(5490702); + } +}; + } // namespace mongo::window_function -- cgit v1.2.1