SERVER-54238 Implement removable $push and $addToSet window functions

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