Skip to content

Commit ce06ac4

Browse files
fotiniAlvanakiEvergreen Agent
fotiniAlvanaki
authored and
Evergreen Agent
committed
SERVER-83394 window stage unit tests
1 parent b39883a commit ce06ac4

File tree

2 files changed

+238
-0
lines changed

2 files changed

+238
-0
lines changed

‎src/mongo/db/query/SConscript

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -520,6 +520,7 @@ env.CppUnitTest(
520520
"view_response_formatter_test.cpp",
521521
'map_reduce_output_format_test.cpp',
522522
'max_time_ms_parser_test.cpp',
523+
"sbe_stage_builder_window_function_test.cpp",
523524
],
524525
LIBDEPS=[
525526
"$BUILD_DIR/mongo/db/auth/authmocks",

‎src/mongo/db/query/sbe_stage_builder_window_function_test.cpp

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,237 @@
1+
/**
2+
* Copyright (C) 2020-present MongoDB, Inc.
3+
*
4+
* This program is free software: you can redistribute it and/or modify
5+
* it under the terms of the Server Side Public License, version 1,
6+
* as published by MongoDB, Inc.
7+
*
8+
* This program is distributed in the hope that it will be useful,
9+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
10+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11+
* Server Side Public License for more details.
12+
*
13+
* You should have received a copy of the Server Side Public License
14+
* along with this program. If not, see
15+
* <http://www.mongodb.com/licensing/server-side-public-license>.
16+
*
17+
* As a special exception, the copyright holders give permission to link the
18+
* code of portions of this program with the OpenSSL library under certain
19+
* conditions as described in each individual source file and distribute
20+
* linked combinations including the program with the OpenSSL library. You
21+
* must comply with the Server Side Public License in all respects for
22+
* all of the code used other than as permitted herein. If you modify file(s)
23+
* with this exception, you may extend this exception to your version of the
24+
* file(s), but you are not obligated to do so. If you do not wish to do so,
25+
* delete this exception statement from your version. If you delete this
26+
* exception statement from all source files in the program, then also delete
27+
* it in the license file.
28+
*/
29+
30+
#include <cstddef>
31+
#include <cstdint>
32+
#include <memory>
33+
#include <tuple>
34+
#include <vector>
35+
36+
#include "mongo/base/string_data.h"
37+
#include "mongo/db/exec/sbe/values/slot.h"
38+
#include "mongo/db/exec/sbe/values/value.h"
39+
#include "mongo/db/pipeline/document_source.h"
40+
#include "mongo/db/pipeline/document_source_set_window_fields.h"
41+
#include "mongo/db/pipeline/expression_context_for_test.h"
42+
#include "mongo/db/query/collation/collator_interface_mock.h"
43+
#include "mongo/db/query/sbe_stage_builder_helpers.h"
44+
#include "mongo/db/query/sbe_stage_builder_test_fixture.h"
45+
#include "mongo/unittest/assert.h"
46+
#include "mongo/unittest/framework.h"
47+
48+
namespace mongo {
49+
50+
51+
class SBESetWindowFieldsTest : public SbeStageBuilderTestFixture {
52+
public:
53+
boost::intrusive_ptr<DocumentSource> createSetWindowFieldsDocumentSource(
54+
boost::intrusive_ptr<ExpressionContext> expCtx, const BSONObj& windowSpec) {
55+
ASSERT(expCtx) << "expCtx must not be null";
56+
BSONObj spec = BSON("$_internalSetWindowFields" << windowSpec);
57+
58+
auto docSrc =
59+
DocumentSourceInternalSetWindowFields::createFromBson(spec.firstElement(), expCtx);
60+
docSrc->optimize();
61+
62+
return docSrc;
63+
}
64+
65+
std::pair<std::unique_ptr<QuerySolution>,
66+
boost::intrusive_ptr<DocumentSourceInternalSetWindowFields>>
67+
makeSetWindowFieldsQuerySolution(boost::intrusive_ptr<ExpressionContext> expCtx,
68+
BSONObj spec,
69+
std::vector<BSONArray> inputDocs) {
70+
71+
auto docSrcA = createSetWindowFieldsDocumentSource(expCtx, spec);
72+
auto docSrc = dynamic_cast<DocumentSourceInternalSetWindowFields*>(docSrcA.get());
73+
ASSERT(docSrc != nullptr);
74+
75+
// Constructs a QuerySolution consisting of a WindowNode on top of a VirtualScanNode.
76+
auto virtScanNode = std::make_unique<VirtualScanNode>(
77+
inputDocs, VirtualScanNode::ScanType::kCollScan, false /*hasRecordId*/);
78+
79+
auto windowNode = std::make_unique<WindowNode>(std::move(virtScanNode),
80+
docSrc->getPartitionBy(),
81+
docSrc->getSortBy(),
82+
docSrc->getOutputFields());
83+
84+
// Makes a QuerySolution from the root window node.
85+
return {makeQuerySolution(std::move(windowNode)), docSrc};
86+
}
87+
88+
std::pair<sbe::value::TypeTags, sbe::value::Value> getSetWindowFieldsResults(
89+
BSONObj windowSpec,
90+
std::vector<BSONArray> inputDocs,
91+
std::unique_ptr<CollatorInterface> collator = nullptr) {
92+
// Makes a QuerySolution for SetWindowFieldsNode over VirtualScanNode.
93+
auto [querySolution, windowNode] =
94+
makeSetWindowFieldsQuerySolution(make_intrusive<ExpressionContextForTest>(),
95+
std::move(windowSpec),
96+
std::move(inputDocs));
97+
98+
// Translates the QuerySolution tree to a sbe::PlanStage tree.
99+
auto [resultSlots, stage, data, _] = buildPlanStage(
100+
std::move(querySolution), false /*hasRecordId*/, nullptr, std::move(collator));
101+
ASSERT_EQ(resultSlots.size(), 1);
102+
103+
auto resultAccessors = prepareTree(&data.env.ctx, stage.get(), resultSlots[0]);
104+
return getAllResults(stage.get(), &resultAccessors[0]);
105+
}
106+
107+
void runSetWindowFieldsTest(StringData windowSpec,
108+
std::vector<BSONArray> inputDocs,
109+
const mongo::BSONArray& expectedValue,
110+
std::unique_ptr<CollatorInterface> collator = nullptr) {
111+
auto [resultsTag, resultsVal] = getSetWindowFieldsResults(
112+
fromjson(windowSpec.rawData()), inputDocs, std::move(collator));
113+
sbe::value::ValueGuard resultGuard{resultsTag, resultsVal};
114+
115+
auto [expectedTag, expectedVal] = stage_builder::makeValue(expectedValue);
116+
sbe::value::ValueGuard expectedGuard{expectedTag, expectedVal};
117+
118+
ASSERT_TRUE(
119+
PlanStageTestFixture::valueEquals(resultsTag, resultsVal, expectedTag, expectedVal))
120+
<< "expected: " << std::make_pair(expectedTag, expectedVal)
121+
<< " but got: " << std::make_pair(resultsTag, resultsVal);
122+
}
123+
};
124+
125+
TEST_F(SBESetWindowFieldsTest, FirstTestPositiveWindow) {
126+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
127+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
128+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
129+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
130+
runSetWindowFieldsTest(
131+
R"({sortBy: {a: 1}, output: {first: {$first: '$b', window: {documents: [1, 2]} }}})",
132+
docs,
133+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "first" << 3)
134+
<< BSON("a" << 2 << "b" << 3 << "first" << 5)
135+
<< BSON("a" << 3 << "b" << 5 << "first" << 7)
136+
<< BSON("a" << 4 << "b" << 7 << "first" << BSONNULL)));
137+
}
138+
139+
TEST_F(SBESetWindowFieldsTest, FirstTestConstantValuePositiveWindow) {
140+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
141+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
142+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
143+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
144+
runSetWindowFieldsTest(
145+
R"({sortBy: {a: 1}, output: {first: {$first: 1000, window: {documents: [1, 2]} }}})",
146+
docs,
147+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "first" << 1000)
148+
<< BSON("a" << 2 << "b" << 3 << "first" << 1000)
149+
<< BSON("a" << 3 << "b" << 5 << "first" << 1000)
150+
<< BSON("a" << 4 << "b" << 7 << "first" << BSONNULL)));
151+
}
152+
153+
TEST_F(SBESetWindowFieldsTest, FirstTestNegativeWindow) {
154+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
155+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
156+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
157+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
158+
runSetWindowFieldsTest(
159+
R"({sortBy: {a: 1}, output: {first: {$first: '$b', window: {documents: [-2, -1]} }}})",
160+
docs,
161+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "first" << BSONNULL)
162+
<< BSON("a" << 2 << "b" << 3 << "first" << 1)
163+
<< BSON("a" << 3 << "b" << 5 << "first" << 1)
164+
<< BSON("a" << 4 << "b" << 7 << "first" << 3)));
165+
}
166+
167+
TEST_F(SBESetWindowFieldsTest, FirstTestConstantValueNegativeWindow) {
168+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
169+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
170+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
171+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
172+
runSetWindowFieldsTest(
173+
R"({sortBy: {a: 1}, output: {first: {$first: 1000, window: {documents: [-2, -1]} }}})",
174+
docs,
175+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "first" << BSONNULL)
176+
<< BSON("a" << 2 << "b" << 3 << "first" << 1000)
177+
<< BSON("a" << 3 << "b" << 5 << "first" << 1000)
178+
<< BSON("a" << 4 << "b" << 7 << "first" << 1000)));
179+
}
180+
181+
TEST_F(SBESetWindowFieldsTest, LastTestPositiveWindow) {
182+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
183+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
184+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
185+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
186+
runSetWindowFieldsTest(
187+
R"({sortBy: {a: 1}, output: {last: {$last: '$b', window: {documents: [1, 2]} }}})",
188+
docs,
189+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "last" << 5)
190+
<< BSON("a" << 2 << "b" << 3 << "last" << 7)
191+
<< BSON("a" << 3 << "b" << 5 << "last" << 7)
192+
<< BSON("a" << 4 << "b" << 7 << "last" << BSONNULL)));
193+
}
194+
195+
TEST_F(SBESetWindowFieldsTest, LastTestConstantValuePositiveWindow) {
196+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
197+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
198+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
199+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
200+
runSetWindowFieldsTest(
201+
R"({sortBy: {a: 1}, output: {last: {$last: 1000, window: {documents: [1, 2]} }}})",
202+
docs,
203+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "last" << 1000)
204+
<< BSON("a" << 2 << "b" << 3 << "last" << 1000)
205+
<< BSON("a" << 3 << "b" << 5 << "last" << 1000)
206+
<< BSON("a" << 4 << "b" << 7 << "last" << BSONNULL)));
207+
}
208+
209+
TEST_F(SBESetWindowFieldsTest, LastTestNegativeWindow) {
210+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
211+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
212+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
213+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
214+
runSetWindowFieldsTest(
215+
R"({sortBy: {a: 1}, output: {last: {$last: '$b', window: {documents: [-2, -1]} }}})",
216+
docs,
217+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "last" << BSONNULL)
218+
<< BSON("a" << 2 << "b" << 3 << "last" << 1)
219+
<< BSON("a" << 3 << "b" << 5 << "last" << 3)
220+
<< BSON("a" << 4 << "b" << 7 << "last" << 5)));
221+
}
222+
223+
TEST_F(SBESetWindowFieldsTest, LastTestConstantValueNegativeWindow) {
224+
auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 1)),
225+
BSON_ARRAY(BSON("a" << 2 << "b" << 3)),
226+
BSON_ARRAY(BSON("a" << 3 << "b" << 5)),
227+
BSON_ARRAY(BSON("a" << 4 << "b" << 7))};
228+
runSetWindowFieldsTest(
229+
R"({sortBy: {a: 1}, output: {last: {$last: 1000, window: {documents: [-2, -1]} }}})",
230+
docs,
231+
BSON_ARRAY(BSON("a" << 1 << "b" << 1 << "last" << BSONNULL)
232+
<< BSON("a" << 2 << "b" << 3 << "last" << 1000)
233+
<< BSON("a" << 3 << "b" << 5 << "last" << 1000)
234+
<< BSON("a" << 4 << "b" << 7 << "last" << 1000)));
235+
}
236+
237+
} // namespace mongo

0 commit comments

Comments
 (0)