diff options
Diffstat (limited to 'src/mongo/db/pipeline')
4 files changed, 384 insertions, 1 deletions
diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript index 806a5d01718..7213e8b7a81 100644 --- a/src/mongo/db/pipeline/SConscript +++ b/src/mongo/db/pipeline/SConscript @@ -427,8 +427,10 @@ env.CppUnitTest( 'skip_and_limit_test.cpp', 'tee_buffer_test.cpp', 'window_function/partition_iterator_test.cpp', + 'window_function/window_function_add_to_set_test.cpp', 'window_function/window_function_exec_test.cpp', 'window_function/window_function_min_max_test.cpp', + 'window_function/window_function_push_test.cpp', ], LIBDEPS=[ '$BUILD_DIR/mongo/base', diff --git a/src/mongo/db/pipeline/window_function/window_function.h b/src/mongo/db/pipeline/window_function/window_function.h index ad348eafa84..b8feba5a2ef 100644 --- a/src/mongo/db/pipeline/window_function/window_function.h +++ b/src/mongo/db/pipeline/window_function/window_function.h @@ -99,7 +99,6 @@ public: MONGO_UNREACHABLE_TASSERT(5371401); } - protected: // Holds all the values in the window, in order, with constant-time access to both ends. ValueMultiset _values; @@ -107,4 +106,104 @@ protected: using WindowFunctionMin = WindowFunctionMinMax<AccumulatorMinMax::Sense::kMin>; using WindowFunctionMax = WindowFunctionMinMax<AccumulatorMinMax::Sense::kMax>; +class WindowFunctionAddToSet final : public WindowFunctionState { +public: + static inline const Value kDefault = Value{std::vector<Value>()}; + + static std::unique_ptr<WindowFunctionState> create(ExpressionContext* const expCtx) { + return std::make_unique<WindowFunctionAddToSet>(expCtx); + } + + explicit WindowFunctionAddToSet(ExpressionContext* const expCtx) + : WindowFunctionState(expCtx), + _values(_expCtx->getValueComparator().makeOrderedValueMultiset()) {} + + void add(Value value) override { + _values.insert(std::move(value)); + } + + /** + * This should only remove the first/lowest element in the window. + */ + void remove(Value value) override { + auto iter = _values.find(std::move(value)); + tassert( + 5423800, "Can't remove from an empty WindowFunctionAddToSet", iter != _values.end()); + _values.erase(iter); + } + + void reset() override { + _values.clear(); + } + + Value getValue() const override { + std::vector<Value> output; + if (_values.empty()) + return kDefault; + for (auto it = _values.begin(); it != _values.end(); it = _values.upper_bound(*it)) { + output.push_back(*it); + } + + return Value(output); + } + +private: + ValueMultiset _values; +}; + +class WindowFunctionPush final : public WindowFunctionState { +public: + using ValueListConstIterator = std::list<Value>::const_iterator; + + static inline const Value kDefault = Value{std::vector<Value>()}; + + static std::unique_ptr<WindowFunctionState> create(ExpressionContext* const expCtx) { + return std::make_unique<WindowFunctionPush>(expCtx); + } + + explicit WindowFunctionPush(ExpressionContext* const expCtx) + : WindowFunctionState(expCtx), + _values( + _expCtx->getValueComparator().makeOrderedValueMultimap<ValueListConstIterator>()) {} + + void add(Value value) override { + _list.emplace_back(std::move(value)); + auto iter = std::prev(_list.end()); + _values.insert({*iter, iter}); + } + + /** + * This should only remove the first/lowest element in the window. + */ + void remove(Value value) override { + // The order of the key-value pairs whose keys compare equivalent is the order of insertion + // and does not change in std::multimap. So find() / erase() will remove the oldest equal + // element, which is what we want, to satisfy "remove() undoes add() when called in FIFO + // order". + auto iter = _values.find(std::move(value)); + tassert(5423801, "Can't remove from an empty WindowFunctionPush", iter != _values.end()); + // Erase the element from both '_values' and '_list'. + _list.erase(iter->second); + _values.erase(iter); + } + + void reset() override { + _values.clear(); + _list.clear(); + } + + Value getValue() const override { + std::vector<Value> output; + if (_values.empty()) + return kDefault; + + return Value{std::vector<Value>(_list.begin(), _list.end())}; + } + +private: + ValueMultimap<ValueListConstIterator> _values; + // std::list makes sure that the order of the elements in the returned array is the order of + // insertion. + std::list<Value> _list; +}; } // namespace mongo diff --git a/src/mongo/db/pipeline/window_function/window_function_add_to_set_test.cpp b/src/mongo/db/pipeline/window_function/window_function_add_to_set_test.cpp new file mode 100644 index 00000000000..6e4d8545240 --- /dev/null +++ b/src/mongo/db/pipeline/window_function/window_function_add_to_set_test.cpp @@ -0,0 +1,129 @@ +/** + * 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 "mongo/platform/basic.h" + +#include "mongo/db/exec/document_value/document_value_test_util.h" +#include "mongo/db/pipeline/aggregation_context_fixture.h" +#include "mongo/db/pipeline/window_function/window_function.h" +#include "mongo/unittest/unittest.h" + +namespace mongo { +namespace { + +class WindowFunctionAddToSetTest : public AggregationContextFixture { +public: + WindowFunctionAddToSetTest() : expCtx(getExpCtx()), addToSet(expCtx.get()) {} + + void addValuesToWindow(const std::vector<Value>& values) { + for (auto val : values) + addToSet.add(val); + } + + boost::intrusive_ptr<ExpressionContext> expCtx; + WindowFunctionAddToSet addToSet; +}; + +TEST_F(WindowFunctionAddToSetTest, EmptyWindowReturnsDefault) { + ASSERT_VALUE_EQ(addToSet.getValue(), Value{std::vector<Value>()}); + + addToSet.add(Value{1}); + addToSet.remove(Value{1}); + ASSERT_VALUE_EQ(addToSet.getValue(), Value{std::vector<Value>()}); +} + +TEST_F(WindowFunctionAddToSetTest, SingletonWindowShouldReturnAVector) { + addToSet.add(Value{5}); + ASSERT_VALUE_EQ(addToSet.getValue(), Value{std::vector<Value>{Value{5}}}); + + addToSet.reset(); + addToSet.add(Value{std::string("str")}); + ASSERT_VALUE_EQ(addToSet.getValue(), Value{std::vector<Value>{Value{std::string("str")}}}); +} + +TEST_F(WindowFunctionAddToSetTest, ComplexWindow) { + addValuesToWindow({Value{2}, Value{5}, Value{std::string("str")}, Value{BSONObj()}}); + + std::vector<Value> expected = {Value{2}, Value{5}, Value{std::string("str")}, Value{BSONObj()}}; + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionAddToSetTest, Removal) { + addValuesToWindow({Value{1}, Value{2}}); + addToSet.remove(Value{1}); + + std::vector<Value> expected = {Value{2}}; + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionAddToSetTest, NotAllowDuplicates) { + addValuesToWindow({Value{1}, Value{1}, Value{2}, Value{3}}); + + // $addToSet window function returns a vector of values that are in a set excluding duplicates. + std::vector<Value> expected = {Value{1}, Value{2}, Value{3}}; + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); + + addToSet.remove(Value{1}); + // Have to remove element '1' twice to remove it from the set. + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); + + addToSet.remove(Value{1}); + expected = {Value{2}, Value{3}}; + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionAddToSetTest, NestedArraysShouldReturnNestedArrays) { + Value vecValue{std::vector<Value>{{Value{1}, Value{2}, Value{3}}}}; + addValuesToWindow({vecValue, vecValue, vecValue}); + + std::vector<Value> expected = {vecValue}; + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionAddToSetTest, IdenticalDocsAndDocsWithDifferentFieldsOrder) { + Value doc1{Document({{"a", 1}, {"b", 2}, {"c", 3}})}; + Value doc2{Document({{"a", 1}, {"c", 3}, {"b", 2}})}; + Value doc3{Document({{"a", 1}, {"b", 2}, {"c", 3}})}; + Value doc4{Document({{"a", 3}, {"b", 3}, {"c", 3}})}; + + addValuesToWindow({doc1, doc2, doc3, doc4}); + + std::vector<Value> expected = {doc1, doc2, doc4}; // doc1 and doc3 are identical. + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); + + addToSet.remove(doc1); + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); + + addToSet.remove(doc2); + expected = {doc1, doc4}; + ASSERT_VALUE_EQ(addToSet.getValue(), Value{expected}); +} + +} // namespace +} // namespace mongo diff --git a/src/mongo/db/pipeline/window_function/window_function_push_test.cpp b/src/mongo/db/pipeline/window_function/window_function_push_test.cpp new file mode 100644 index 00000000000..e96e62c0f26 --- /dev/null +++ b/src/mongo/db/pipeline/window_function/window_function_push_test.cpp @@ -0,0 +1,153 @@ +/** + * 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 "mongo/platform/basic.h" + +#include "mongo/db/exec/document_value/document_value_test_util.h" +#include "mongo/db/pipeline/aggregation_context_fixture.h" +#include "mongo/db/pipeline/window_function/window_function.h" +#include "mongo/unittest/unittest.h" + +namespace mongo { +namespace { + +class WindowFunctionPushTest : public AggregationContextFixture { +public: + WindowFunctionPushTest() : expCtx(getExpCtx()), push(expCtx.get()) {} + + void addValuesToWindow(const std::vector<Value>& values) { + for (auto val : values) + push.add(val); + } + + boost::intrusive_ptr<ExpressionContext> expCtx; + WindowFunctionPush push; +}; + +TEST_F(WindowFunctionPushTest, EmptyWindowReturnsDefault) { + ASSERT_VALUE_EQ(push.getValue(), Value{std::vector<Value>()}); + + push.add(Value{1}); + push.remove(Value{1}); + ASSERT_VALUE_EQ(push.getValue(), Value{std::vector<Value>()}); +} + +TEST_F(WindowFunctionPushTest, SingletonWindowShouldReturnAVector) { + push.add(Value{5}); + ASSERT_VALUE_EQ(push.getValue(), Value(std::vector<Value>{Value{5}})); + + push.reset(); + push.add(Value{std::string("str")}); + ASSERT_VALUE_EQ(push.getValue(), Value(std::vector<Value>{Value{std::string("str")}})); +} + +TEST_F(WindowFunctionPushTest, ComplexWindowReservesInsertionOrder) { + addValuesToWindow({Value{std::string("str")}, Value{5}, Value{2}, Value{BSONObj()}}); + + std::vector<Value> expected = {Value{std::string("str")}, Value{5}, Value{2}, Value{BSONObj()}}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionPushTest, Removal) { + addValuesToWindow({Value{1}, Value{2}}); + push.remove(Value{1}); + + std::vector<Value> expected = {Value{2}}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionPushTest, AllowsDuplicates) { + addValuesToWindow({Value{1}, Value{1}, Value{2}, Value{3}}); + + // $push allows duplicates in the array returned by the window function. + std::vector<Value> expected = {Value{1}, Value{1}, Value{2}, Value{3}}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); + + // remove() value '1' once only removes one instance in the array. + push.remove(Value{1}); + expected = {Value{1}, Value{2}, Value{3}}; + + push.remove(Value{1}); + expected = {Value{2}, Value{3}}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionPushTest, ShouldReserveInsertionOrder) { + addValuesToWindow({Value{1234}, Value{123}, Value{321}, Value{10000}}); + + std::vector<Value> expected = {Value{1234}, Value{123}, Value{321}, Value{10000}}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionPushTest, RemovalDoesNotAffectOrder) { + addValuesToWindow({Value{1234}, Value{123}, Value{321}, Value{10000}}); + + std::vector<Value> expected = {Value{1234}, Value{123}, Value{321}, Value{10000}}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); + + push.remove(Value{1234}); + expected = {Value{123}, Value{321}, Value{10000}}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionPushTest, NestedArraysShouldReturnNestedArrays) { + Value vecValue{std::vector<Value>{{Value{1}, Value{2}, Value{3}}}}; + addValuesToWindow({vecValue, vecValue, vecValue}); + + std::vector<Value> expected = {vecValue, vecValue, vecValue}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); + + push.remove(vecValue); + push.remove(vecValue); + expected = {vecValue}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); +} + +TEST_F(WindowFunctionPushTest, IdenticalDocsAndDocsWithDifferentFieldsOrder) { + Value doc1{Document({{"a", 1}, {"b", 2}, {"c", 3}})}; + Value doc2{Document({{"a", 1}, {"b", 2}, {"c", 3}})}; + Value doc3{Document({{"a", 1}, {"c", 3}, {"b", 2}})}; + Value doc4{Document({{"a", 3}, {"b", 3}, {"c", 3}})}; + + addValuesToWindow({doc1, doc2, doc3, doc4}); + + std::vector<Value> expected = {doc1, doc2, doc3, doc4}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); + + push.remove(doc1); + expected = {doc2, doc3, doc4}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); + + push.remove(doc2); + expected = {doc3, doc4}; + ASSERT_VALUE_EQ(push.getValue(), Value{expected}); +} + +} // namespace +} // namespace mongo |