SERVER-55455 Pass variables state through window function input evaluation

4 files changed, 47 insertions, 31 deletions

diff --git a/src/mongo/db/pipeline/window_function/window_function_exec_non_removable.h b/src/mongo/db/pipeline/window_function/window_function_exec_non_removable.h
index d00e668ef53..bb7418d66c8 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec_non_removable.h
+++ b/src/mongo/db/pipeline/window_function/window_function_exec_non_removable.h
@@ -79,7 +79,8 @@ public:
             }();
 
             if (auto doc = (this->_iter)[upperIndex])
-                _function->process(_input->evaluate(*doc, nullptr), false);
+                _function->process(
+                    _input->evaluate(*doc, &_input->getExpressionContext()->variables), false);
             else {
                 // Upper bound is out of range, but may be because it's off of the end of the
                 // partition. For instance, for bounds [unbounded, -1] we won't be able to
@@ -122,7 +123,8 @@ private:
         _initialized = true;
         for (int i = 0; needMore(i); i++) {
             if (auto doc = (this->_iter)[i]) {
-                _function->process(_input->evaluate(*doc, nullptr), false);
+                _function->process(
+                    _input->evaluate(*doc, &_input->getExpressionContext()->variables), false);
             } else {
                 // Already reached the end of partition for the first value to compute.
                 break;
diff --git a/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp b/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp
index 62738d664ab..14b9ee2ac44 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_exec_non_removable_test.cpp
@@ -72,33 +72,6 @@ private:
     std::unique_ptr<PartitionIterator> _iter;
 };
 
-class WindowFunctionExecRemovableDocumentTest : public AggregationContextFixture {
-public:
-    WindowFunctionExecRemovableDocumentTest() {}
-    WindowFunctionExecRemovableDocument createForFieldPath(
-        std::deque<DocumentSource::GetNextResult> docs,
-        const std::string& inputPath,
-        WindowBounds::DocumentBased bounds) {
-        _docSource = DocumentSourceMock::createForTest(std::move(docs), getExpCtx());
-        _iter =
-            std::make_unique<PartitionIterator>(getExpCtx().get(), _docSource.get(), boost::none);
-        auto input = ExpressionFieldPath::parse(
-            getExpCtx().get(), inputPath, getExpCtx()->variablesParseState);
-        std::unique_ptr<WindowFunctionState> maxFunc =
-            std::make_unique<WindowFunctionMax>(getExpCtx().get());
-        return WindowFunctionExecRemovableDocument(
-            _iter.get(), std::move(input), std::move(maxFunc), bounds);
-    }
-
-    auto advanceIterator() {
-        return _iter->advance();
-    }
-
-private:
-    boost::intrusive_ptr<DocumentSourceMock> _docSource;
-    std::unique_ptr<PartitionIterator> _iter;
-};
-
 TEST_F(WindowFunctionExecNonRemovableTest, AccumulateOnlyWithoutLookahead) {
     const auto docs = std::deque<DocumentSource::GetNextResult>{
         Document{{"a", 1}}, Document{{"a", 2}}, Document{{"a", 3}}};
@@ -178,5 +151,25 @@ TEST_F(WindowFunctionExecNonRemovableTest, FullPartitionWindow) {
     ASSERT_VALUE_EQ(Value(6), mgr.getNext());
 }
 
+TEST_F(WindowFunctionExecNonRemovableTest, InputExpressionAllowedToCreateVariables) {
+    const auto docs = std::deque<DocumentSource::GetNextResult>{
+        Document{{"a", 1}}, Document{{"a", 2}}, Document{{"a", 3}}};
+    auto docSource = DocumentSourceMock::createForTest(std::move(docs), getExpCtx());
+    auto iter =
+        std::make_unique<PartitionIterator>(getExpCtx().get(), docSource.get(), boost::none);
+    auto filterBSON =
+        fromjson("{$filter: {input: [1, 2, 3], as: 'num', cond: {$gte: ['$$num', 2]}}}");
+    auto input = ExpressionFilter::parse(
+        getExpCtx().get(), filterBSON.firstElement(), getExpCtx()->variablesParseState);
+    auto exec = WindowFunctionExecNonRemovable<AccumulatorState>(
+        iter.get(), std::move(input), AccumulatorFirst::create(getExpCtx().get()), 1);
+    // The input is a constant [2, 3] for each document.
+    ASSERT_VALUE_EQ(Value({Value(2), Value(3)}), exec.getNext());
+    iter->advance();
+    ASSERT_VALUE_EQ(Value({Value(2), Value(3)}), exec.getNext());
+    iter->advance();
+    ASSERT_VALUE_EQ(Value({Value(2), Value(3)}), exec.getNext());
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/window_function/window_function_exec_removable_document.cpp b/src/mongo/db/pipeline/window_function/window_function_exec_removable_document.cpp
index 1bab909467b..a459d295a8e 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec_removable_document.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_exec_removable_document.cpp
@@ -68,7 +68,7 @@ void WindowFunctionExecRemovableDocument::initialize() {
     for (int i = lowerBoundForInit; !_upperBound || i <= _upperBound.get(); ++i) {
         // If this is false, we're over the end of the partition.
         if (auto doc = (this->_iter)[i]) {
-            addValue(_input->evaluate(*doc, nullptr));
+            addValue(_input->evaluate(*doc, &_input->getExpressionContext()->variables));
         } else {
             break;
         }
@@ -80,7 +80,7 @@ void WindowFunctionExecRemovableDocument::processDocumentsToUpperBound() {
     if (_upperBound) {
         // If this is false, we're over the end of the partition.
         if (auto doc = (this->_iter)[_upperBound.get()]) {
-            addValue(_input->evaluate(*doc, nullptr));
+            addValue(_input->evaluate(*doc, &_input->getExpressionContext()->variables));
         }
     }
 }
diff --git a/src/mongo/db/pipeline/window_function/window_function_exec_removable_test.cpp b/src/mongo/db/pipeline/window_function/window_function_exec_removable_test.cpp
index f98e26c526d..5eae82ab741 100644
--- a/src/mongo/db/pipeline/window_function/window_function_exec_removable_test.cpp
+++ b/src/mongo/db/pipeline/window_function/window_function_exec_removable_test.cpp
@@ -306,5 +306,26 @@ TEST_F(WindowFunctionExecRemovableDocumentTest, CanResetFunction) {
     }
 }
 
+TEST_F(WindowFunctionExecRemovableDocumentTest, InputExpressionAllowedToCreateVariables) {
+    const auto docs = std::deque<DocumentSource::GetNextResult>{
+        Document{{"a", 1}}, Document{{"a", 2}}, Document{{"a", 3}}};
+    auto docSource = DocumentSourceMock::createForTest(std::move(docs), getExpCtx());
+    auto iter =
+        std::make_unique<PartitionIterator>(getExpCtx().get(), docSource.get(), boost::none);
+    auto filterBSON =
+        fromjson("{$filter: {input: [1, 2, 3], as: 'num', cond: {$gte: ['$$num', 2]}}}");
+    auto input = ExpressionFilter::parse(
+        getExpCtx().get(), filterBSON.firstElement(), getExpCtx()->variablesParseState);
+    auto maxFunc = std::make_unique<WindowFunctionMax>(getExpCtx().get());
+    auto mgr = WindowFunctionExecRemovableDocument(
+        iter.get(), std::move(input), std::move(maxFunc), WindowBounds::DocumentBased{-1, 0});
+    // The input is a constant [2, 3] for each document.
+    ASSERT_VALUE_EQ(Value({Value(2), Value(3)}), mgr.getNext());
+    iter->advance();
+    ASSERT_VALUE_EQ(Value({Value(2), Value(3)}), mgr.getNext());
+    iter->advance();
+    ASSERT_VALUE_EQ(Value({Value(2), Value(3)}), mgr.getNext());
+}
+
 }  // namespace
 }  // namespace mongo