summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorEric Cox <eric.cox@mongodb.com>2020-10-23 14:44:01 +0000
committerEvergreen Agent <no-reply@evergreen.mongodb.com>2020-11-11 14:26:55 +0000
commitb43f6fefe1ac3e941fd55d5452a1ee21e7ff0ae6 (patch)
tree79f280087ed6024bab99f8315eb5723459b1adce
parent8e7b31355bfce777f46d77c960257b1801a714b0 (diff)
downloadmongo-b43f6fefe1ac3e941fd55d5452a1ee21e7ff0ae6.tar.gz
SERVER-51788 Implement infrastructure for testing SBE stage builder
-rw-r--r--src/mongo/db/SConscript2
-rw-r--r--src/mongo/db/exec/sbe/SConscript15
-rw-r--r--src/mongo/db/exec/sbe/sbe_filter_test.cpp26
-rw-r--r--src/mongo/db/exec/sbe/sbe_limit_skip_test.cpp10
-rw-r--r--src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp115
-rw-r--r--src/mongo/db/exec/sbe/sbe_plan_stage_test.h71
-rw-r--r--src/mongo/db/exec/sbe/sbe_sort_test.cpp4
-rw-r--r--src/mongo/db/exec/sbe/sbe_sorted_merge_test.cpp19
-rw-r--r--src/mongo/db/exec/sbe/sbe_unique_test.cpp20
-rw-r--r--src/mongo/db/query/SConscript18
-rw-r--r--src/mongo/db/query/classic_stage_builder.cpp24
-rw-r--r--src/mongo/db/query/classic_stage_builder_test.cpp145
-rw-r--r--src/mongo/db/query/query_solution.cpp23
-rw-r--r--src/mongo/db/query/query_solution.h45
-rw-r--r--src/mongo/db/query/sbe_stage_builder.cpp33
-rw-r--r--src/mongo/db/query/sbe_stage_builder.h1
-rw-r--r--src/mongo/db/query/sbe_stage_builder_helpers.cpp70
-rw-r--r--src/mongo/db/query/sbe_stage_builder_helpers.h26
-rw-r--r--src/mongo/db/query/sbe_stage_builder_test.cpp110
-rw-r--r--src/mongo/db/query/sbe_stage_builder_test_fixture.cpp74
-rw-r--r--src/mongo/db/query/sbe_stage_builder_test_fixture.h79
-rw-r--r--src/mongo/db/query/stage_types.h4
22 files changed, 768 insertions, 166 deletions
diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 9866cb342e1..698477f284b 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -1197,7 +1197,6 @@ env.Library(
'query/sbe_stage_builder_filter.cpp',
'query/sbe_stage_builder_index_scan.cpp',
'query/sbe_stage_builder_projection.cpp',
- 'query/sbe_stage_builder_helpers.cpp',
'query/sbe_sub_planner.cpp',
'query/stage_builder_util.cpp',
'query/wildcard_multikey_paths.cpp',
@@ -1241,6 +1240,7 @@ env.Library(
'query/plan_yield_policy',
'query/query_common',
'query/query_planner',
+ 'query/sbe_stage_builder_helpers',
'repl/repl_coordinator_interface',
's/sharding_api_d',
'shared_request_handling',
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index f2d9ba6a87a..ec71185af3c 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -98,6 +98,18 @@ env.Library(
]
)
+env.Library(
+ target='sbe_plan_stage_test',
+ source=[
+ 'sbe_plan_stage_test.cpp',
+ ],
+ LIBDEPS=[
+ '$BUILD_DIR/mongo/db/query/sbe_stage_builder_helpers',
+ '$BUILD_DIR/mongo/unittest/unittest',
+ 'query_sbe',
+ ]
+ )
+
env.CppUnitTest(
target='db_sbe_test',
source=[
@@ -120,7 +132,6 @@ env.CppUnitTest(
'sbe_limit_skip_test.cpp',
'sbe_math_builtins_test.cpp',
'sbe_numeric_convert_test.cpp',
- 'sbe_plan_stage_test.cpp',
'sbe_sort_test.cpp',
'sbe_sorted_merge_test.cpp',
'sbe_test.cpp',
@@ -130,7 +141,7 @@ env.CppUnitTest(
LIBDEPS=[
'$BUILD_DIR/mongo/db/concurrency/lock_manager',
'$BUILD_DIR/mongo/db/service_context_test_fixture',
- '$BUILD_DIR/mongo/unittest/unittest',
'query_sbe_parser',
+ 'sbe_plan_stage_test',
],
)
diff --git a/src/mongo/db/exec/sbe/sbe_filter_test.cpp b/src/mongo/db/exec/sbe/sbe_filter_test.cpp
index 033d436cdf9..eb7c7c05214 100644
--- a/src/mongo/db/exec/sbe/sbe_filter_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_filter_test.cpp
@@ -43,8 +43,8 @@ namespace mongo::sbe {
using FilterStageTest = PlanStageTestFixture;
TEST_F(FilterStageTest, ConstantFilterAlwaysTrueTest) {
- auto [inputTag, inputVal] =
- makeValue(BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
+ auto [inputTag, inputVal] = stage_builder::makeValue(
+ BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
value::ValueGuard inputGuard{inputTag, inputVal};
auto [expectedTag, expectedVal] = value::copyValue(inputTag, inputVal);
@@ -66,8 +66,8 @@ TEST_F(FilterStageTest, ConstantFilterAlwaysTrueTest) {
}
TEST_F(FilterStageTest, ConstantFilterAlwaysFalseTest) {
- auto [inputTag, inputVal] =
- makeValue(BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
+ auto [inputTag, inputVal] = stage_builder::makeValue(
+ BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
value::ValueGuard inputGuard{inputTag, inputVal};
auto [expectedTag, expectedVal] = value::makeNewArray();
@@ -89,8 +89,8 @@ TEST_F(FilterStageTest, ConstantFilterAlwaysFalseTest) {
}
TEST_F(FilterStageTest, FilterAlwaysTrueTest) {
- auto [inputTag, inputVal] =
- makeValue(BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
+ auto [inputTag, inputVal] = stage_builder::makeValue(
+ BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
value::ValueGuard inputGuard{inputTag, inputVal};
auto [expectedTag, expectedVal] = value::copyValue(inputTag, inputVal);
@@ -112,8 +112,8 @@ TEST_F(FilterStageTest, FilterAlwaysTrueTest) {
}
TEST_F(FilterStageTest, FilterAlwaysFalseTest) {
- auto [inputTag, inputVal] =
- makeValue(BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
+ auto [inputTag, inputVal] = stage_builder::makeValue(
+ BSON_ARRAY(12LL << "yar" << BSON_ARRAY(2.5) << 7.5 << BSON("foo" << 23)));
value::ValueGuard inputGuard{inputTag, inputVal};
auto [expectedTag, expectedVal] = value::makeNewArray();
@@ -137,11 +137,11 @@ TEST_F(FilterStageTest, FilterAlwaysFalseTest) {
TEST_F(FilterStageTest, FilterIsNumberTest) {
using namespace std::literals;
- auto [inputTag, inputVal] =
- makeValue(BSON_ARRAY(12LL << "42" << BSON_ARRAY(2.5) << 7.5 << BSON("34" << 56)));
+ auto [inputTag, inputVal] = stage_builder::makeValue(
+ BSON_ARRAY(12LL << "42" << BSON_ARRAY(2.5) << 7.5 << BSON("34" << 56)));
value::ValueGuard inputGuard{inputTag, inputVal};
- auto [expectedTag, expectedVal] = makeValue(BSON_ARRAY(12LL << 7.5));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(BSON_ARRAY(12LL << 7.5));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto makeStageFn = [](value::SlotId scanSlot, std::unique_ptr<PlanStage> scanStage) {
@@ -160,13 +160,13 @@ TEST_F(FilterStageTest, FilterIsNumberTest) {
}
TEST_F(FilterStageTest, FilterLessThanTest) {
- auto [inputTag, inputVal] = makeValue(BSON_ARRAY(
+ auto [inputTag, inputVal] = stage_builder::makeValue(BSON_ARRAY(
BSON_ARRAY(2.8 << 3) << BSON_ARRAY(7LL << 5.0) << BSON_ARRAY(4LL << 4.3)
<< BSON_ARRAY(8 << 8) << BSON_ARRAY("1" << 2) << BSON_ARRAY(1 << "2")
<< BSON_ARRAY(4.9 << 5) << BSON_ARRAY(6.0 << BSON_ARRAY(11.0))));
value::ValueGuard inputGuard{inputTag, inputVal};
- auto [expectedTag, expectedVal] = makeValue(
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(
BSON_ARRAY(BSON_ARRAY(2.8 << 3) << BSON_ARRAY(4LL << 4.3) << BSON_ARRAY(4.9 << 5)));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
diff --git a/src/mongo/db/exec/sbe/sbe_limit_skip_test.cpp b/src/mongo/db/exec/sbe/sbe_limit_skip_test.cpp
index 023fdee286c..69eec391e80 100644
--- a/src/mongo/db/exec/sbe/sbe_limit_skip_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_limit_skip_test.cpp
@@ -41,11 +41,13 @@ namespace mongo::sbe {
using LimitSkipStageTest = PlanStageTestFixture;
TEST_F(LimitSkipStageTest, LimitSimpleTest) {
+ auto ctx = makeCompileCtx();
+
// Make a "limit 1000" stage.
auto limit = makeS<LimitSkipStage>(
makeS<CoScanStage>(kEmptyPlanNodeId), 1000, boost::none, kEmptyPlanNodeId);
- prepareTree(limit.get());
+ prepareTree(ctx.get(), limit.get());
// Verify that `limit` produces at least 1000 values.
for (int i = 0; i < 1000; ++i) {
@@ -57,6 +59,8 @@ TEST_F(LimitSkipStageTest, LimitSimpleTest) {
}
TEST_F(LimitSkipStageTest, LimitSkipSimpleTest) {
+ auto ctx = makeCompileCtx();
+
// Make an input array containing 64-integers 0 thru 999, inclusive.
auto [inputTag, inputVal] = value::makeNewArray();
value::ValueGuard inputGuard{inputTag, inputVal};
@@ -68,10 +72,10 @@ TEST_F(LimitSkipStageTest, LimitSkipSimpleTest) {
// Make a "limit 200 skip 300" stage.
inputGuard.reset();
- auto [scanSlot, scanStage] = generateMockScan(inputTag, inputVal);
+ auto [scanSlot, scanStage] = generateVirtualScan(inputTag, inputVal);
auto limit = makeS<LimitSkipStage>(std::move(scanStage), 200, 300, kEmptyPlanNodeId);
- auto resultAccessor = prepareTree(limit.get(), scanSlot);
+ auto resultAccessor = prepareTree(ctx.get(), limit.get(), scanSlot);
// Verify that `limit` produces exactly 300 thru 499, inclusive.
for (i = 0; i < 200; ++i) {
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
index c634160f3d3..ef0bc49822e 100644
--- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
@@ -39,108 +39,39 @@
namespace mongo::sbe {
-std::pair<value::TypeTags, value::Value> PlanStageTestFixture::makeValue(const BSONArray& ba) {
- int numBytes = ba.objsize();
- uint8_t* data = new uint8_t[numBytes];
- memcpy(data, reinterpret_cast<const uint8_t*>(ba.objdata()), numBytes);
- return {value::TypeTags::bsonArray, value::bitcastFrom<uint8_t*>(data)};
-}
-
-std::pair<value::TypeTags, value::Value> PlanStageTestFixture::makeValue(const BSONObj& bo) {
- int numBytes = bo.objsize();
- uint8_t* data = new uint8_t[numBytes];
- memcpy(data, reinterpret_cast<const uint8_t*>(bo.objdata()), numBytes);
- return {value::TypeTags::bsonObject, value::bitcastFrom<uint8_t*>(data)};
-}
-
-std::pair<value::SlotId, std::unique_ptr<PlanStage>> PlanStageTestFixture::generateMockScan(
- value::TypeTags arrTag, value::Value arrVal) {
- // The value passed in must be an array.
- invariant(value::isArray(arrTag));
-
- // Make an EConstant expression for the array.
- auto arrayExpression = makeE<EConstant>(arrTag, arrVal);
-
- // Build the unwind/project/limit/coscan subtree.
- auto projectSlot = generateSlotId();
- auto unwindSlot = generateSlotId();
- auto unwind = makeS<UnwindStage>(
- makeProjectStage(
- makeS<LimitSkipStage>(
- makeS<CoScanStage>(kEmptyPlanNodeId), 1, boost::none, kEmptyPlanNodeId),
- kEmptyPlanNodeId,
- projectSlot,
- std::move(arrayExpression)),
- projectSlot,
- unwindSlot,
- generateSlotId(), // We don't need an index slot but must to provide it.
- false, // Don't preserve null and empty arrays.
- kEmptyPlanNodeId);
-
- // Return the UnwindStage and its output slot. The UnwindStage can be used as an input
- // to other PlanStages.
- return {unwindSlot, std::move(unwind)};
-}
-
-std::pair<value::SlotVector, std::unique_ptr<PlanStage>>
-PlanStageTestFixture::generateMockScanMulti(int32_t numSlots,
- value::TypeTags arrTag,
- value::Value arrVal) {
- using namespace std::literals;
-
- invariant(numSlots >= 1);
-
- // Generate a mock scan with a single output slot.
- auto [scanSlot, scanStage] = generateMockScan(arrTag, arrVal);
-
- // Create a ProjectStage that will read the data from `scanStage` and split it up
- // across multiple output slots.
- value::SlotVector projectSlots;
- value::SlotMap<std::unique_ptr<EExpression>> projections;
- for (int32_t i = 0; i < numSlots; ++i) {
- projectSlots.emplace_back(generateSlotId());
- projections.emplace(
- projectSlots.back(),
- makeE<EFunction>("getElement"sv,
- makeEs(makeE<EVariable>(scanSlot),
- makeE<EConstant>(value::TypeTags::NumberInt32,
- value::bitcastFrom<int32_t>(i)))));
- }
-
- return {std::move(projectSlots),
- makeS<ProjectStage>(std::move(scanStage), std::move(projections), kEmptyPlanNodeId)};
-}
-
-std::pair<value::SlotId, std::unique_ptr<PlanStage>> PlanStageTestFixture::generateMockScan(
+std::pair<value::SlotId, std::unique_ptr<PlanStage>> PlanStageTestFixture::generateVirtualScan(
const BSONArray& array) {
- auto [arrTag, arrVal] = makeValue(array);
- return generateMockScan(arrTag, arrVal);
+ auto [arrTag, arrVal] = stage_builder::makeValue(array);
+ return generateVirtualScan(arrTag, arrVal);
}
std::pair<value::SlotVector, std::unique_ptr<PlanStage>>
-PlanStageTestFixture::generateMockScanMulti(int32_t numSlots, const BSONArray& array) {
- auto [arrTag, arrVal] = makeValue(array);
- return generateMockScanMulti(numSlots, arrTag, arrVal);
+PlanStageTestFixture::generateVirtualScanMulti(int32_t numSlots, const BSONArray& array) {
+ auto [arrTag, arrVal] = stage_builder::makeValue(array);
+ return generateVirtualScanMulti(numSlots, arrTag, arrVal);
}
-void PlanStageTestFixture::prepareTree(PlanStage* root) {
- root->prepare(*_compileCtx);
+void PlanStageTestFixture::prepareTree(CompileCtx* ctx, PlanStage* root) {
+ root->prepare(*ctx);
root->attachFromOperationContext(opCtx());
root->open(false);
}
-value::SlotAccessor* PlanStageTestFixture::prepareTree(PlanStage* root, value::SlotId slot) {
- prepareTree(root);
- return root->getAccessor(*_compileCtx, slot);
+value::SlotAccessor* PlanStageTestFixture::prepareTree(CompileCtx* ctx,
+ PlanStage* root,
+ value::SlotId slot) {
+ prepareTree(ctx, root);
+ return root->getAccessor(*ctx, slot);
}
-std::vector<value::SlotAccessor*> PlanStageTestFixture::prepareTree(PlanStage* root,
+std::vector<value::SlotAccessor*> PlanStageTestFixture::prepareTree(CompileCtx* ctx,
+ PlanStage* root,
value::SlotVector slots) {
std::vector<value::SlotAccessor*> slotAccessors;
- prepareTree(root);
+ prepareTree(ctx, root);
for (auto slot : slots) {
- slotAccessors.emplace_back(root->getAccessor(*_compileCtx, slot));
+ slotAccessors.emplace_back(root->getAccessor(*ctx, slot));
}
return slotAccessors;
}
@@ -193,18 +124,20 @@ void PlanStageTestFixture::runTest(value::TypeTags inputTag,
value::TypeTags expectedTag,
value::Value expectedVal,
const MakeStageFn<value::SlotId>& makeStage) {
+ auto ctx = makeCompileCtx();
+
// Set up a ValueGuard to ensure `expected` gets released.
value::ValueGuard expectedGuard{expectedTag, expectedVal};
// Generate a mock scan from `input` with a single output slot.
- auto [scanSlot, scanStage] = generateMockScan(inputTag, inputVal);
+ auto [scanSlot, scanStage] = generateVirtualScan(inputTag, inputVal);
// Call the `makeStage` callback to create the PlanStage that we want to test, passing in
// the mock scan subtree and its output slot.
auto [outputSlot, stage] = makeStage(scanSlot, std::move(scanStage));
// Prepare the tree and get the SlotAccessor for the output slot.
- auto resultAccessor = prepareTree(stage.get(), outputSlot);
+ auto resultAccessor = prepareTree(ctx.get(), stage.get(), outputSlot);
// Get all the results produced by the PlanStage we want to test.
auto [resultsTag, resultsVal] = getAllResults(stage.get(), resultAccessor);
@@ -220,18 +153,20 @@ void PlanStageTestFixture::runTestMulti(int32_t numInputSlots,
value::TypeTags expectedTag,
value::Value expectedVal,
const MakeStageFn<value::SlotVector>& makeStageMulti) {
+ auto ctx = makeCompileCtx();
+
// Set up a ValueGuard to ensure `expected` gets released.
value::ValueGuard expectedGuard{expectedTag, expectedVal};
// Generate a mock scan from `input` with multiple output slots.
- auto [scanSlots, scanStage] = generateMockScanMulti(numInputSlots, inputTag, inputVal);
+ auto [scanSlots, scanStage] = generateVirtualScanMulti(numInputSlots, inputTag, inputVal);
// Call the `makeStageMulti` callback to create the PlanStage that we want to test, passing
// in the mock scan subtree and its output slots.
auto [outputSlots, stage] = makeStageMulti(scanSlots, std::move(scanStage));
// Prepare the tree and get the SlotAccessors for the output slots.
- auto resultAccessors = prepareTree(stage.get(), outputSlots);
+ auto resultAccessors = prepareTree(ctx.get(), stage.get(), outputSlots);
// Get all the results produced by the PlanStage we want to test.
auto [resultsTag, resultsVal] = getAllResultsMulti(stage.get(), resultAccessors);
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
index c8d74595c35..b1cef9d319a 100644
--- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
+++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
@@ -40,6 +40,8 @@
#include "mongo/db/exec/sbe/values/bson.h"
#include "mongo/db/exec/sbe/values/slot.h"
#include "mongo/db/exec/sbe/values/value.h"
+#include "mongo/db/query/sbe_stage_builder.h"
+#include "mongo/db/query/sbe_stage_builder_helpers.h"
#include "mongo/db/service_context_test_fixture.h"
#include "mongo/unittest/unittest.h"
@@ -53,14 +55,14 @@ using MakeStageFn = std::function<std::pair<T, std::unique_ptr<PlanStage>>(
* PlanStageTestFixture is a unittest framework for testing sbe::PlanStages.
*
* To facilitate writing unittests for PlanStages, PlanStageTestFixture sets up an OperationContext
- * and a CompileCtx and offers a number of methods to help unittest writers. From the perspective a
- * unittest writer, the most important methods in the PlanStageTestFixture class are prepareTree(),
- * runTest(), and runTestMulti(). Each unittest should directly call only one of these methods once.
+ * and offers a number of methods to help unittest writers. From the perspective a unittest writer,
+ * the most important methods in the PlanStageTestFixture class are prepareTree(), runTest(), and
+ * runTestMulti(). Each unittest should directly call only one of these methods once.
*
* For unittests where you need more control and flexibility, calling prepareTree() directly is
- * the way to go. prepareTree() takes the root stage of a PlanStage tree and 0 or more SlotIds as
- * parameters. When invoked, prepareTree() calls prepare() on the root stage (passing in the
- * CompileCtx), attaches the OperationContext to the root stage, calls open() on the root stage,
+ * the way to go. prepareTree() takes a CompileCtx, the root stage of a PlanStage tree and 0 or more
+ * SlotIds as parameters. When invoked, prepareTree() calls prepare() on the root stage (passing in
+ * the CompileCtx), attaches the OperationContext to the root stage, calls open() on the root stage,
* and then returns the SlotAccessors corresponding to the specified SlotIds. For a given unittest
* that calls prepareTree() directly, you can think of the unittest as having two parts: (1) the
* part before prepareTree(); and (2) the part after prepareTree(). The first part of the test
@@ -68,8 +70,9 @@ using MakeStageFn = std::function<std::pair<T, std::unique_ptr<PlanStage>>(
* The second part of the test (after prepareTree()) should drive the execution of the PlanStage
* tree (by calling getNext() on the root stage one or more times) and verify that the PlanStage
* tree behaves as expected. During the first part before prepareTree(), it's common to use
- * generateMockScan() or generateMockScanMulti() which provide an easy way to build a PlanStage
- * subtree that streams out the contents of an SBE array (mimicking a real collection scan).
+ * generateVirtualScan() or generateVirtualScanMulti() which provide an easy way to build a
+ * PlanStage subtree that streams out the contents of an SBE array (mimicking a real collection
+ * scan).
*
* For unittests where you just need to stream the contents of an input array to a PlanStage and
* compare the values produced against an "expected output" array, runTest() or runTestMulti() are
@@ -85,11 +88,9 @@ public:
ServiceContextTest::setUp();
_opCtx = makeOperationContext();
_slotIdGenerator.reset(new value::SlotIdGenerator());
- _compileCtx.reset(new CompileCtx(std::make_unique<RuntimeEnvironment>()));
}
void tearDown() override {
- _compileCtx.reset();
_slotIdGenerator.reset();
_opCtx.reset();
ServiceContextTest::tearDown();
@@ -103,8 +104,11 @@ public:
return _slotIdGenerator->generate();
}
- CompileCtx* compileCtx() {
- return _compileCtx.get();
+ /**
+ * Makes a new CompileCtx suitable for preparing an sbe::PlanStage tree.
+ */
+ std::unique_ptr<CompileCtx> makeCompileCtx() {
+ return std::make_unique<CompileCtx>(std::make_unique<RuntimeEnvironment>());
}
/**
@@ -119,29 +123,19 @@ public:
}
/**
- * Converts a BSONArray to an SBE Array. Caller owns the SBE Array returned. This method
- * does not assume ownership of the BSONArray.
- */
- std::pair<value::TypeTags, value::Value> makeValue(const BSONArray& ba);
-
- /**
- * Converts a BSONObj to an SBE Object. Caller owns the SBE Object returned. This method
- * does not assume ownership of the BSONObj.
- */
- std::pair<value::TypeTags, value::Value> makeValue(const BSONObj& bo);
-
- /**
* This method takes an SBE array and returns an output slot and a unwind/project/limit/coscan
* subtree that streams out the elements of the array one at a time via the output slot over a
* series of calls to getNext(), mimicking the output of a collection scan or an index scan.
*
* Note that this method assumes ownership of the SBE Array being passed in.
*/
- std::pair<value::SlotId, std::unique_ptr<PlanStage>> generateMockScan(value::TypeTags arrTag,
- value::Value arrVal);
+ std::pair<value::SlotId, std::unique_ptr<PlanStage>> generateVirtualScan(value::TypeTags arrTag,
+ value::Value arrVal) {
+ return stage_builder::generateVirtualScan(_slotIdGenerator.get(), arrTag, arrVal);
+ };
/**
- * This method is similar to generateMockScan(), except that the subtree returned outputs to
+ * This method is similar to generateVirtualScan(), except that the subtree returned outputs to
* multiple slots instead of a single slot. `numSlots` specifies the number of output slots.
* `array` is expected to be an array of subarrays. Each subarray is expected to have exactly
* `numSlots` elements, where the value at index 0 corresponds to output slot 0, the value at
@@ -151,37 +145,43 @@ public:
*
* Note that this method assumes ownership of the SBE Array being passed in.
*/
- std::pair<value::SlotVector, std::unique_ptr<PlanStage>> generateMockScanMulti(
- int32_t numSlots, value::TypeTags arrTag, value::Value arrVal);
+ std::pair<value::SlotVector, std::unique_ptr<PlanStage>> generateVirtualScanMulti(
+ int32_t numSlots, value::TypeTags arrTag, value::Value arrVal) {
+ return stage_builder::generateVirtualScanMulti(
+ _slotIdGenerator.get(), numSlots, arrTag, arrVal);
+ };
/**
* Make a mock scan from an BSON array. This method does NOT assume ownership of the BSONArray
* passed in.
*/
- std::pair<value::SlotId, std::unique_ptr<PlanStage>> generateMockScan(const BSONArray& array);
+ std::pair<value::SlotId, std::unique_ptr<PlanStage>> generateVirtualScan(
+ const BSONArray& array);
/**
* Make a mock scan with multiple output slots from an BSON array. This method does NOT assume
* ownership of the BSONArray passed in.
*/
- std::pair<value::SlotVector, std::unique_ptr<PlanStage>> generateMockScanMulti(
+ std::pair<value::SlotVector, std::unique_ptr<PlanStage>> generateVirtualScanMulti(
int32_t numSlots, const BSONArray& array);
/**
- * Prepares the tree of PlanStages given by `root` and returns the SlotAccessor* for `slot`.
+ * Prepares the tree of PlanStages given by `root`.
*/
- void prepareTree(PlanStage* root);
+ void prepareTree(CompileCtx* ctx, PlanStage* root);
/**
* Prepares the tree of PlanStages given by `root` and returns the SlotAccessor* for `slot`.
*/
- value::SlotAccessor* prepareTree(PlanStage* root, value::SlotId slot);
+ value::SlotAccessor* prepareTree(CompileCtx* ctx, PlanStage* root, value::SlotId slot);
/**
* Prepares the tree of PlanStages given by `root` and returns the SlotAccessor*'s for
* the specified slots.
*/
- std::vector<value::SlotAccessor*> prepareTree(PlanStage* root, value::SlotVector slots);
+ std::vector<value::SlotAccessor*> prepareTree(CompileCtx* ctx,
+ PlanStage* root,
+ value::SlotVector slots);
/**
* This method repeatedly calls getNext() on the specified PlanStage, stores all the values
@@ -242,7 +242,6 @@ public:
private:
ServiceContext::UniqueOperationContext _opCtx;
std::unique_ptr<value::SlotIdGenerator> _slotIdGenerator;
- std::unique_ptr<CompileCtx> _compileCtx;
};
} // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/sbe_sort_test.cpp b/src/mongo/db/exec/sbe/sbe_sort_test.cpp
index 659d5bae9a6..d471d37f413 100644
--- a/src/mongo/db/exec/sbe/sbe_sort_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_sort_test.cpp
@@ -39,12 +39,12 @@ namespace mongo::sbe {
using SortStageTest = PlanStageTestFixture;
TEST_F(SortStageTest, SortNumbersTest) {
- auto [inputTag, inputVal] = makeValue(
+ auto [inputTag, inputVal] = stage_builder::makeValue(
BSON_ARRAY(BSON_ARRAY(12LL << "A") << BSON_ARRAY(2.5 << "B") << BSON_ARRAY(7 << "C")
<< BSON_ARRAY(Decimal128(4) << "D")));
value::ValueGuard inputGuard{inputTag, inputVal};
- auto [expectedTag, expectedVal] = makeValue(
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(
BSON_ARRAY(BSON_ARRAY(2.5 << "B") << BSON_ARRAY(Decimal128(4) << "D")
<< BSON_ARRAY(7 << "C") << BSON_ARRAY(12LL << "A")));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
diff --git a/src/mongo/db/exec/sbe/sbe_sorted_merge_test.cpp b/src/mongo/db/exec/sbe/sbe_sorted_merge_test.cpp
index df167bfbafa..37cad8d4532 100644
--- a/src/mongo/db/exec/sbe/sbe_sorted_merge_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_sorted_merge_test.cpp
@@ -59,9 +59,9 @@ public:
std::vector<std::unique_ptr<PlanStage>> inputScans;
std::vector<value::SlotVector> inputSlots;
for (auto&& array : inputStreams) {
- auto [tag, val] = makeValue(array);
+ auto [tag, val] = stage_builder::makeValue(array);
- auto [slotVector, scan] = generateMockScanMulti(numSlots, tag, val);
+ auto [slotVector, scan] = generateVirtualScanMulti(numSlots, tag, val);
inputScans.push_back(std::move(scan));
inputSlots.push_back(std::move(slotVector));
}
@@ -87,7 +87,8 @@ public:
outputVals,
kEmptyPlanNodeId);
- auto resultAccessors = prepareTree(sortedMerge.get(), outputVals);
+ auto ctx = makeCompileCtx();
+ auto resultAccessors = prepareTree(ctx.get(), sortedMerge.get(), outputVals);
return {resultAccessors, std::move(sortedMerge)};
}
@@ -116,7 +117,7 @@ TEST_F(SortedMergeStageTest, TwoChildrenBasicAscending) {
BSON_ARRAY(BSON_ARRAY(2 << 2) << BSON_ARRAY(4 << 4))},
std::vector<value::SortDirection>{value::SortDirection::Ascending});
- auto [expectedTag, expectedVal] = makeValue(BSON_ARRAY(1 << 2 << 3 << 4));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(BSON_ARRAY(1 << 2 << 3 << 4));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto [resultsTag, resultsVal] = getAllResults(sortedMerge.get(), resultAccessors[0]);
@@ -131,7 +132,7 @@ TEST_F(SortedMergeStageTest, TwoChildrenBasicDescending) {
BSON_ARRAY(BSON_ARRAY(4 << 4) << BSON_ARRAY(2 << 2) << BSON_ARRAY(1 << 1))},
std::vector<value::SortDirection>{value::SortDirection::Descending});
- auto [expectedTag, expectedVal] = makeValue(BSON_ARRAY(5 << 4 << 3 << 2 << 1));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(BSON_ARRAY(5 << 4 << 3 << 2 << 1));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto [resultsTag, resultsVal] = getAllResults(sortedMerge.get(), resultAccessors[0]);
@@ -145,7 +146,7 @@ TEST_F(SortedMergeStageTest, TwoChildrenOneEmpty) {
makeSortedMerge({BSONArray(), BSON_ARRAY(BSON_ARRAY(1 << 1) << BSON_ARRAY(2 << 2))},
std::vector<value::SortDirection>{value::SortDirection::Ascending});
- auto [expectedTag, expectedVal] = makeValue(BSON_ARRAY(1 << 2));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(BSON_ARRAY(1 << 2));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto [resultsTag, resultsVal] = getAllResults(sortedMerge.get(), resultAccessors[0]);
@@ -160,7 +161,7 @@ TEST_F(SortedMergeStageTest, TwoChildrenWithDuplicates) {
BSON_ARRAY(BSON_ARRAY(1 << 1) << BSON_ARRAY(2 << 2))},
std::vector<value::SortDirection>{value::SortDirection::Ascending});
- auto [expectedTag, expectedVal] = makeValue(BSON_ARRAY(1 << 1 << 1 << 2));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(BSON_ARRAY(1 << 1 << 1 << 2));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto [resultsTag, resultsVal] = getAllResults(sortedMerge.get(), resultAccessors[0]);
@@ -179,8 +180,8 @@ TEST_F(SortedMergeStageTest, FiveChildren) {
<< BSON_ARRAY(10 << 10) << BSON_ARRAY(11 << 11) << BSON_ARRAY(12 << 12))},
std::vector<value::SortDirection>{value::SortDirection::Ascending});
- auto [expectedTag, expectedVal] =
- makeValue(BSON_ARRAY(1 << 1 << 1 << 1 << 2 << 3 << 4 << 4 << 5 << 10 << 11 << 12));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(
+ BSON_ARRAY(1 << 1 << 1 << 1 << 2 << 3 << 4 << 4 << 5 << 10 << 11 << 12));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto [resultsTag, resultsVal] = getAllResults(sortedMerge.get(), resultAccessors[0]);
diff --git a/src/mongo/db/exec/sbe/sbe_unique_test.cpp b/src/mongo/db/exec/sbe/sbe_unique_test.cpp
index d1556594df2..08ed98c1723 100644
--- a/src/mongo/db/exec/sbe/sbe_unique_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_unique_test.cpp
@@ -33,6 +33,7 @@
#include "mongo/db/exec/sbe/sbe_plan_stage_test.h"
#include "mongo/db/exec/sbe/stages/unique.h"
+#include "mongo/db/query/sbe_stage_builder_helpers.h"
namespace mongo::sbe {
/**
@@ -42,10 +43,10 @@ namespace mongo::sbe {
using UniqueStageTest = PlanStageTestFixture;
TEST_F(UniqueStageTest, DeduplicatesAndPreservesOrderSimple) {
- auto [inputTag, inputVal] = makeValue(BSON_ARRAY(1 << 2 << 3 << 1));
+ auto [inputTag, inputVal] = stage_builder::makeValue(BSON_ARRAY(1 << 2 << 3 << 1));
value::ValueGuard inputGuard{inputTag, inputVal};
- auto [expectedTag, expectedVal] = makeValue(BSON_ARRAY(1 << 2 << 3));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(BSON_ARRAY(1 << 2 << 3));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto makeStageFn = [](value::SlotId scanSlot, std::unique_ptr<PlanStage> scanStage) {
@@ -61,19 +62,20 @@ TEST_F(UniqueStageTest, DeduplicatesAndPreservesOrderSimple) {
}
TEST_F(UniqueStageTest, DeduplicatesMultipleSlotsInKey) {
- auto [tag, val] = makeValue(BSON_ARRAY(
+ auto [tag, val] = stage_builder::makeValue(BSON_ARRAY(
BSON_ARRAY(1 << 1) << BSON_ARRAY(2 << 2) << BSON_ARRAY(1 << 1) << BSON_ARRAY(3 << 3)));
- auto [scanSlots, scan] = generateMockScanMulti(2, // numSlots
- tag,
- val);
+ auto [scanSlots, scan] = generateVirtualScanMulti(2, // numSlots
+ tag,
+ val);
- auto [expectedTag, expectedVal] =
- makeValue(BSON_ARRAY(BSON_ARRAY(1 << 1) << BSON_ARRAY(2 << 2) << BSON_ARRAY(3 << 3)));
+ auto [expectedTag, expectedVal] = stage_builder::makeValue(
+ BSON_ARRAY(BSON_ARRAY(1 << 1) << BSON_ARRAY(2 << 2) << BSON_ARRAY(3 << 3)));
value::ValueGuard expectedGuard{expectedTag, expectedVal};
auto unique = makeS<UniqueStage>(std::move(scan), scanSlots, kEmptyPlanNodeId);
- auto resultAccessors = prepareTree(unique.get(), scanSlots);
+ auto ctx = makeCompileCtx();
+ auto resultAccessors = prepareTree(ctx.get(), unique.get(), scanSlots);
auto [resultsTag, resultsVal] = getAllResultsMulti(unique.get(), resultAccessors);
value::ValueGuard resultGuard{resultsTag, resultsVal};
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript
index eba0ff73786..b4741d22748 100644
--- a/src/mongo/db/query/SConscript
+++ b/src/mongo/db/query/SConscript
@@ -75,6 +75,18 @@ env.Library(
)
env.Library(
+ target='sbe_stage_builder_helpers',
+ source=[
+ "sbe_stage_builder_helpers.cpp",
+ ],
+ LIBDEPS=[
+ "$BUILD_DIR/mongo/base",
+ '$BUILD_DIR/mongo/db/exec/sbe/query_sbe',
+ "$BUILD_DIR/mongo/db/matcher/expressions",
+ ],
+)
+
+env.Library(
target='projection_ast',
source=[
"projection.cpp",
@@ -281,6 +293,7 @@ env.CppUnitTest(
source=[
"canonical_query_encoder_test.cpp",
"canonical_query_test.cpp",
+ "classic_stage_builder_test.cpp",
"count_command_test.cpp",
"cursor_response_test.cpp",
"explain_options_test.cpp",
@@ -324,17 +337,22 @@ env.CppUnitTest(
"query_request_test.cpp",
"query_settings_test.cpp",
"query_solution_test.cpp",
+ "sbe_stage_builder_test_fixture.cpp",
+ "sbe_stage_builder_test.cpp",
"view_response_formatter_test.cpp",
],
LIBDEPS=[
"$BUILD_DIR/mongo/db/auth/authmocks",
"$BUILD_DIR/mongo/db/concurrency/lock_manager",
+ "$BUILD_DIR/mongo/db/exec/sbe/sbe_plan_stage_test",
"$BUILD_DIR/mongo/db/pipeline/aggregation_request",
"$BUILD_DIR/mongo/db/query_exec",
"$BUILD_DIR/mongo/db/service_context_d",
+ "$BUILD_DIR/mongo/db/service_context_d_test_fixture",
"$BUILD_DIR/mongo/db/service_context_test_fixture",
"$BUILD_DIR/mongo/dbtests/mocklib",
"$BUILD_DIR/mongo/rpc/rpc",
+ "$BUILD_DIR/mongo/util/clock_source_mock",
"collation/collator_factory_mock",
"collation/collator_interface_mock",
"command_request_response",
diff --git a/src/mongo/db/query/classic_stage_builder.cpp b/src/mongo/db/query/classic_stage_builder.cpp
index 39a872951f5..dfe01081c8f 100644
--- a/src/mongo/db/query/classic_stage_builder.cpp
+++ b/src/mongo/db/query/classic_stage_builder.cpp
@@ -53,6 +53,7 @@
#include "mongo/db/exec/merge_sort.h"
#include "mongo/db/exec/or.h"
#include "mongo/db/exec/projection.h"
+#include "mongo/db/exec/queued_data_stage.h"
#include "mongo/db/exec/return_key.h"
#include "mongo/db/exec/shard_filter.h"
#include "mongo/db/exec/skip.h"
@@ -354,6 +355,29 @@ std::unique_ptr<PlanStage> ClassicStageBuilder::build(const QuerySolutionNode* r
case STAGE_EOF: {
return std::make_unique<EOFStage>(expCtx);
}
+ case STAGE_VIRTUAL_SCAN: {
+ const auto* vsn = static_cast<const VirtualScanNode*>(root);
+ invariant(vsn->hasRecordId);
+
+ auto qds = std::make_unique<QueuedDataStage>(expCtx, _ws);
+ for (auto&& arr : vsn->docs) {
+ // The VirtualScanNode should only have a single element that carrys the document
+ // as the QueuedDataStage cannot handle a recordId properly.
+ BSONObjIterator arrIt{arr};
+ invariant(arrIt.more());
+ auto firstElt = arrIt.next();
+ invariant(firstElt.type() == BSONType::Object);
+ invariant(!arrIt.more());
+
+ // Only add the first element to the working set.
+ auto wsID = _ws->allocate();
+ qds->pushBack(wsID);
+ auto* member = _ws->get(wsID);
+ member->keyData.clear();
+ member->doc = {{}, Document{firstElt.embeddedObject()}};
+ }
+ return qds;
+ }
case STAGE_CACHED_PLAN:
case STAGE_COUNT:
case STAGE_DELETE:
diff --git a/src/mongo/db/query/classic_stage_builder_test.cpp b/src/mongo/db/query/classic_stage_builder_test.cpp
new file mode 100644
index 00000000000..5d912b81d3d
--- /dev/null
+++ b/src/mongo/db/query/classic_stage_builder_test.cpp
@@ -0,0 +1,145 @@
+/**
+ * Copyright (C) 2020-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/catalog/collection.h"
+#include "mongo/db/namespace_string.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/query/classic_stage_builder.h"
+#include "mongo/db/query/query_solution.h"
+#include "mongo/db/service_context_d_test_fixture.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/clock_source_mock.h"
+
+namespace mongo {
+
+const static NamespaceString kNss("db.dummy");
+
+class ClassicStageBuilderTest : public ServiceContextMongoDTest {
+public:
+ void setUp() {
+ getServiceContext()->setFastClockSource(std::make_unique<ClockSourceMock>());
+ _opCtx = makeOperationContext();
+ _workingSet = std::make_unique<WorkingSet>();
+ }
+
+ void tearDown() {
+ _opCtx.reset();
+ _workingSet.reset();
+ }
+
+ /**
+ * Converts a 'QuerySolutionNode' to a 'QuerySolution'.
+ */
+ std::unique_ptr<QuerySolution> makeQuerySolution(std::unique_ptr<QuerySolutionNode> root) {
+ auto querySoln = std::make_unique<QuerySolution>();
+ querySoln->setRoot(std::move(root));
+ return querySoln;
+ }
+
+ /**
+ * Builds a PlanStage using the given WorkingSet and QuerySolution.
+ */
+ std::unique_ptr<PlanStage> buildPlanStage(std::unique_ptr<QuerySolution> querySolution) {
+ auto qr = std::make_unique<QueryRequest>(kNss);
+ auto expCtx = make_intrusive<ExpressionContext>(opCtx(), nullptr, kNss);
+ auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr), expCtx);
+ ASSERT_OK(statusWithCQ.getStatus());
+
+ stage_builder::ClassicStageBuilder builder{
+ opCtx(), CollectionPtr::null, *statusWithCQ.getValue(), *querySolution, workingSet()};
+ return builder.build(querySolution->root());
+ }
+
+ /**
+ * A helper to repeatedly call work() until the stage returns a PlanStage::IS_EOF state and
+ * returns the resulting documents as a vector of BSONObj.
+ */
+ std::vector<BSONObj> collectResults(std::unique_ptr<PlanStage> stage) {
+ WorkingSetID id;
+ std::vector<BSONObj> results;
+ auto state = PlanStage::ADVANCED;
+
+ while (state != PlanStage::IS_EOF) {
+ state = stage->work(&id);
+ if (state == PlanStage::ADVANCED) {
+ auto member = workingSet()->get(id);
+ auto doc = member->doc.value().toBson();
+ results.push_back(doc);
+ }
+ }
+ return results;
+ }
+
+ OperationContext* opCtx() {
+ return _opCtx.get();
+ }
+
+ WorkingSet* workingSet() {
+ return _workingSet.get();
+ }
+
+private:
+ ServiceContext::UniqueOperationContext _opCtx;
+ std::unique_ptr<WorkingSet> _workingSet;
+};
+
+
+/**
+ * Verify that a VirtualScanNode can be translated to a QueuedDataStage and produce a data stream.
+ */
+TEST_F(ClassicStageBuilderTest, VirtualScanTranslation) {
+ auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 1 << "b" << 2)),
+ BSON_ARRAY(BSON("a" << 2 << "b" << 2)),
+ BSON_ARRAY(BSON("a" << 3 << "b" << 2))};
+
+ // Construct a QuerySolution consisting of a single VirtualScanNode to test if a stream of
+ // documents can be produced.
+ auto virtScan = std::make_unique<VirtualScanNode>(docs, true);
+
+ // Make a QuerySolution from the root virtual scan node.
+ auto querySolution = makeQuerySolution(std::move(virtScan));
+ ASSERT_EQ(querySolution->root()->nodeId(), 1);
+
+ // Translate the QuerySolution to a classic PlanStage.
+ auto stage = buildPlanStage(std::move(querySolution));
+
+ // Work the stage and collect the results.
+ auto results = collectResults(std::move(stage));
+ ASSERT_EQ(results.size(), 3);
+
+ // Check that the results produced from the translated VirtualScanNode meet expectation.
+ for (size_t i = 0; i < docs.size(); ++i) {
+ BSONObjIterator arrIt{docs[i]};
+ auto firstElt = arrIt.next();
+ ASSERT_BSONOBJ_EQ(firstElt.embeddedObject(), results[i]);
+ }
+}
+} // namespace mongo
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index f0146a19e3f..fe6b33d591f 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -234,6 +234,29 @@ QuerySolutionNode* CollectionScanNode::clone() const {
}
//
+// VirtualScanNode
+//
+
+VirtualScanNode::VirtualScanNode(std::vector<BSONArray> docs, bool hasRecordId)
+ : docs(std::move(docs)), hasRecordId(hasRecordId) {}
+
+void VirtualScanNode::appendToString(str::stream* ss, int indent) const {
+ addIndent(ss, indent);
+ *ss << "VIRTUAL_SCAN\n";
+ addIndent(ss, indent + 1);
+ *ss << "nDocuments = " << docs.size();
+ addIndent(ss, indent + 1);
+ *ss << "hasRecordId = " << hasRecordId;
+ addCommon(ss, indent);
+}
+
+QuerySolutionNode* VirtualScanNode::clone() const {
+ auto copy = new VirtualScanNode(docs, this->hasRecordId);
+ cloneBaseData(copy);
+ return copy;
+}
+
+//
// AndHashNode
//
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index 7801a3ea18e..ef5888a401f 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -482,6 +482,51 @@ struct CollectionScanNode : public QuerySolutionNodeWithSortSet {
bool stopApplyingFilterAfterFirstMatch = false;
};
+/**
+ * A VirtualScanNode is similar to a collection or an index scan except that it doesn't depend on an
+ * underlying storage implementation. It can be used to represent a virtual
+ * collection or an index scan in memory by using a backing vector of BSONArray.
+ */
+struct VirtualScanNode : public QuerySolutionNodeWithSortSet {
+ VirtualScanNode(std::vector<BSONArray> docs, bool hasRecordId);
+ virtual ~VirtualScanNode() {}
+
+ virtual StageType getType() const {
+ return STAGE_VIRTUAL_SCAN;
+ }
+
+ virtual void appendToString(str::stream* ss, int indent) const;
+
+ bool fetched() const {
+ return true;
+ }
+ FieldAvailability getFieldAvailability(const std::string& field) const {
+ return FieldAvailability::kFullyProvided;
+ }
+ bool sortedByDiskLoc() const {
+ return false;
+ }
+
+ QuerySolutionNode* clone() const;
+
+ // A representation of a collection's documents. Here we use a BSONArray so metadata like a
+ // RecordId can be stored alongside of the main document payload. The format of the data in
+ // BSONArray is entirely up to a client of this node, but if this node is to be used for
+ // consumption downstream by stage builder implementations it must conform to the format
+ // expected by those stage builders. That expected contract depends on the hasRecordId flag. If
+ // the hasRecordId flag is 'false' the BSONArray will have a single element that is a BSONObj
+ // representation of a document being produced from this node. If 'hasRecordId' is true, then
+ // each BSONArray in docs will carry a RecordId in the zeroth position of the array and a
+ // BSONObj in the first position of the array.
+ std::vector<BSONArray> docs;
+
+ // A flag to indicate the format of the BSONArray document payload in the above vector, docs. If
+ // hasRecordId is set to true, then both a RecordId and a BSONObj document are stored in that
+ // order for every BSONArray in docs. Otherwise, the RecordId is omitted and the BSONArray will
+ // only carry a BSONObj document.
+ bool hasRecordId;
+};
+
struct AndHashNode : public QuerySolutionNode {
AndHashNode();
virtual ~AndHashNode();
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index cc17cccf974..238fd569146 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -52,6 +52,7 @@
#include "mongo/db/index/fts_access_method.h"
#include "mongo/db/query/sbe_stage_builder_coll_scan.h"
#include "mongo/db/query/sbe_stage_builder_filter.h"
+#include "mongo/db/query/sbe_stage_builder_helpers.h"
#include "mongo/db/query/sbe_stage_builder_index_scan.h"
#include "mongo/db/query/sbe_stage_builder_projection.h"
#include "mongo/db/query/util/make_data_structure.h"
@@ -101,6 +102,34 @@ std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildCollScan(
return std::move(stage);
}
+std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildVirtualScan(
+ const QuerySolutionNode* root) {
+ auto vsn = static_cast<const VirtualScanNode*>(root);
+
+ auto [inputTag, inputVal] = sbe::value::makeNewArray();
+ sbe::value::ValueGuard inputGuard{inputTag, inputVal};
+ auto inputView = sbe::value::getArrayView(inputVal);
+
+ for (auto& doc : vsn->docs) {
+ auto [tag, val] = makeValue(doc);
+ inputView->push_back(tag, val);
+ }
+
+ inputGuard.reset();
+ auto [scanSlots, scanStage] =
+ generateVirtualScanMulti(&_slotIdGenerator, vsn->hasRecordId ? 2 : 1, inputTag, inputVal);
+
+ if (vsn->hasRecordId) {
+ invariant(scanSlots.size() == 2);
+ _data.recordIdSlot = scanSlots[0];
+ _data.resultSlot = scanSlots[1];
+ } else {
+ invariant(scanSlots.size() == 1);
+ _data.resultSlot = scanSlots[0];
+ }
+ return std::move(scanStage);
+}
+
std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildIndexScan(
const QuerySolutionNode* root) {
auto ixn = static_cast<const IndexScanNode*>(root);
@@ -772,6 +801,7 @@ std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::build(const QuerySolution
SlotBasedStageBuilder&, const QuerySolutionNode* root)>>
kStageBuilders = {
{STAGE_COLLSCAN, std::mem_fn(&SlotBasedStageBuilder::buildCollScan)},
+ {STAGE_VIRTUAL_SCAN, std::mem_fn(&SlotBasedStageBuilder::buildVirtualScan)},
{STAGE_IXSCAN, std::mem_fn(&SlotBasedStageBuilder::buildIndexScan)},
{STAGE_FETCH, std::mem_fn(&SlotBasedStageBuilder::buildFetch)},
{STAGE_LIMIT, std::mem_fn(&SlotBasedStageBuilder::buildLimit)},
@@ -786,7 +816,8 @@ std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::build(const QuerySolution
{STAGE_TEXT, &SlotBasedStageBuilder::buildText},
{STAGE_RETURN_KEY, &SlotBasedStageBuilder::buildReturnKey},
{STAGE_EOF, &SlotBasedStageBuilder::buildEof},
- {STAGE_SORT_MERGE, &SlotBasedStageBuilder::buildSortMerge}};
+ {STAGE_SORT_MERGE, &SlotBasedStageBuilder::buildSortMerge},
+ {STAGE_VIRTUAL_SCAN, &SlotBasedStageBuilder::buildVirtualScan}};
uassert(4822884,
str::stream() << "Can't build exec tree for node: " << root->toString(),
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index f2a21960069..ab6da07ceb0 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -111,6 +111,7 @@ public:
private:
std::unique_ptr<sbe::PlanStage> buildCollScan(const QuerySolutionNode* root);
+ std::unique_ptr<sbe::PlanStage> buildVirtualScan(const QuerySolutionNode* root);
std::unique_ptr<sbe::PlanStage> buildIndexScan(const QuerySolutionNode* root);
std::unique_ptr<sbe::PlanStage> buildFetch(const QuerySolutionNode* root);
std::unique_ptr<sbe::PlanStage> buildLimit(const QuerySolutionNode* root);
diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.cpp b/src/mongo/db/query/sbe_stage_builder_helpers.cpp
index f0ca2f22ecf..77a58529d35 100644
--- a/src/mongo/db/query/sbe_stage_builder_helpers.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_helpers.cpp
@@ -35,8 +35,10 @@
#include "mongo/db/exec/sbe/stages/co_scan.h"
#include "mongo/db/exec/sbe/stages/limit_skip.h"
#include "mongo/db/exec/sbe/stages/loop_join.h"
+#include "mongo/db/exec/sbe/stages/project.h"
#include "mongo/db/exec/sbe/stages/traverse.h"
#include "mongo/db/exec/sbe/stages/union.h"
+#include "mongo/db/exec/sbe/stages/unwind.h"
#include "mongo/db/matcher/matcher_type_set.h"
namespace mongo::stage_builder {
@@ -316,4 +318,72 @@ EvalExprStagePair generateShortCircuitingLogicalOp(sbe::EPrimBinary::Op logicOp,
return generateSingleResultUnion(std::move(branches), branchFn, planNodeId, slotIdGenerator);
}
+std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateVirtualScan(
+ sbe::value::SlotIdGenerator* slotIdGenerator,
+ sbe::value::TypeTags arrTag,
+ sbe::value::Value arrVal) {
+ // The value passed in must be an array.
+ invariant(sbe::value::isArray(arrTag));
+
+ // Make an EConstant expression for the array.
+ auto arrayExpression = sbe::makeE<sbe::EConstant>(arrTag, arrVal);
+
+ // Build the unwind/project/limit/coscan subtree.
+ auto projectSlot = slotIdGenerator->generate();
+ auto unwindSlot = slotIdGenerator->generate();
+ auto unwind = sbe::makeS<sbe::UnwindStage>(
+ sbe::makeProjectStage(makeLimitCoScanTree(kEmptyPlanNodeId, 1),
+ kEmptyPlanNodeId,
+ projectSlot,
+ std::move(arrayExpression)),
+ projectSlot,
+ unwindSlot,
+ slotIdGenerator->generate(), // We don't need an index slot but must to provide it.
+ false, // Don't preserve null and empty arrays.
+ kEmptyPlanNodeId);
+
+ // Return the UnwindStage and its output slot. The UnwindStage can be used as an input
+ // to other PlanStages.
+ return {unwindSlot, std::move(unwind)};
+}
+
+std::pair<sbe::value::SlotVector, std::unique_ptr<sbe::PlanStage>> generateVirtualScanMulti(
+ sbe::value::SlotIdGenerator* slotIdGenerator,
+ int numSlots,
+ sbe::value::TypeTags arrTag,
+ sbe::value::Value arrVal) {
+ using namespace std::literals;
+
+ invariant(numSlots >= 1);
+
+ // Generate a mock scan with a single output slot.
+ auto [scanSlot, scanStage] = generateVirtualScan(slotIdGenerator, arrTag, arrVal);
+
+ // Create a ProjectStage that will read the data from 'scanStage' and split it up
+ // across multiple output slots.
+ sbe::value::SlotVector projectSlots;
+ sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projections;
+ for (int32_t i = 0; i < numSlots; ++i) {
+ projectSlots.emplace_back(slotIdGenerator->generate());
+ projections.emplace(
+ projectSlots.back(),
+ sbe::makeE<sbe::EFunction>(
+ "getElement"sv,
+ sbe::makeEs(sbe::makeE<sbe::EVariable>(scanSlot),
+ sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt32,
+ sbe::value::bitcastFrom<int32_t>(i)))));
+ }
+
+ return {std::move(projectSlots),
+ sbe::makeS<sbe::ProjectStage>(
+ std::move(scanStage), std::move(projections), kEmptyPlanNodeId)};
+}
+
+std::pair<sbe::value::TypeTags, sbe::value::Value> makeValue(const BSONArray& ba) {
+ int numBytes = ba.objsize();
+ uint8_t* data = new uint8_t[numBytes];
+ memcpy(data, reinterpret_cast<const uint8_t*>(ba.objdata()), numBytes);
+ return {sbe::value::TypeTags::bsonArray, sbe::value::bitcastFrom<uint8_t*>(data)};
+}
+
} // namespace mongo::stage_builder
diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.h b/src/mongo/db/query/sbe_stage_builder_helpers.h
index bf0e5d27157..6ab645883bd 100644
--- a/src/mongo/db/query/sbe_stage_builder_helpers.h
+++ b/src/mongo/db/query/sbe_stage_builder_helpers.h
@@ -237,4 +237,30 @@ EvalExprStagePair generateShortCircuitingLogicalOp(sbe::EPrimBinary::Op logicOp,
PlanNodeId planNodeId,
sbe::value::SlotIdGenerator* slotIdGenerator);
+/** This helper takes an SBE SlotIdGenerator and an SBE Array and returns an output slot and a
+ * unwind/project/limit/coscan subtree that streams out the elements of the array one at a time via
+ * the output slot over a series of calls to getNext(), mimicking the output of a collection scan or
+ * an index scan. Note that this method assumes ownership of the SBE Array being passed in.
+ */
+std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateVirtualScan(
+ sbe::value::SlotIdGenerator* slotIdGenerator,
+ sbe::value::TypeTags arrTag,
+ sbe::value::Value arrVal);
+
+/**
+ * Make a mock scan with multiple output slots from an BSON array. This method does NOT assume
+ * ownership of the BSONArray passed in.
+ */
+std::pair<sbe::value::SlotVector, std::unique_ptr<sbe::PlanStage>> generateVirtualScanMulti(
+ sbe::value::SlotIdGenerator* slotIdGenerator,
+ int numSlots,
+ sbe::value::TypeTags arrTag,
+ sbe::value::Value arrVal);
+
+/**
+ * Converts a BSONArray to an SBE Array. Caller owns the SBE Array returned. This method does not
+ * assume ownership of the BSONArray.
+ */
+std::pair<sbe::value::TypeTags, sbe::value::Value> makeValue(const BSONArray& ba);
+
} // namespace mongo::stage_builder
diff --git a/src/mongo/db/query/sbe_stage_builder_test.cpp b/src/mongo/db/query/sbe_stage_builder_test.cpp
new file mode 100644
index 00000000000..c830f33bdcb
--- /dev/null
+++ b/src/mongo/db/query/sbe_stage_builder_test.cpp
@@ -0,0 +1,110 @@
+/**
+ * Copyright (C) 2020-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/query/query_solution.h"
+#include "mongo/db/query/sbe_stage_builder_test_fixture.h"
+#include "mongo/unittest/unittest.h"
+
+namespace mongo {
+
+using SBEStageBuilderTest = SBEStageBuilderTestFixture;
+
+TEST_F(SBEStageBuilderTest, TestVirtualScan) {
+ auto docs = std::vector<BSONArray>{BSON_ARRAY(int64_t{0} << BSON("a" << 1 << "b" << 2)),
+ BSON_ARRAY(int64_t{1} << BSON("a" << 2 << "b" << 2)),
+ BSON_ARRAY(int64_t{2} << BSON("a" << 3 << "b" << 2))};
+
+ // Construct a QuerySolution consisting of a single VirtualScanNode to test if a stream of
+ // documents can be produced.
+ auto virtScan = std::make_unique<VirtualScanNode>(docs, true);
+
+ // Make a QuerySolution from the root virtual scan node.
+ auto querySolution = makeQuerySolution(std::move(virtScan));
+ ASSERT_EQ(querySolution->root()->nodeId(), 1);
+
+ // Translate the QuerySolution tree to an sbe::PlanStage.
+ auto [resultSlots, stage, data] = buildPlanStage(std::move(querySolution), true);
+ auto resultAccessors = prepareTree(&data.ctx, stage.get(), resultSlots);
+
+ int64_t index = 0;
+ for (auto st = stage->getNext(); st == sbe::PlanState::ADVANCED; st = stage->getNext()) {
+ // Assert that the recordIDs are what we expect.
+ auto [tag, val] = resultAccessors[0]->getViewOfValue();
+ ASSERT_TRUE(tag == sbe::value::TypeTags::NumberInt64);
+ ASSERT_EQ(index, sbe::value::bitcastTo<int64_t>(val));
+
+ // Assert that the document produced from the stage is what we expect.
+ auto [tagDoc, valDoc] = resultAccessors[1]->getViewOfValue();
+ ASSERT_TRUE(tagDoc == sbe::value::TypeTags::bsonObject);
+ auto bo = BSONObj(sbe::value::bitcastTo<const char*>(valDoc));
+ ASSERT_BSONOBJ_EQ(bo, BSON("a" << ++index << "b" << 2));
+ }
+ ASSERT_EQ(index, 3);
+}
+
+TEST_F(SBEStageBuilderTest, TestLimitOneVirtualScan) {
+ auto docs = std::vector<BSONArray>{BSON_ARRAY(int64_t{0} << BSON("a" << 1 << "b" << 2)),
+ BSON_ARRAY(int64_t{1} << BSON("a" << 2 << "b" << 2)),
+ BSON_ARRAY(int64_t{2} << BSON("a" << 3 << "b" << 2))};
+
+ // Construct a QuerySolution consisting of a root limit node that takes ownership of a
+ // VirtualScanNode.
+ auto virtScan = std::make_unique<VirtualScanNode>(docs, true);
+ auto limitNode = std::make_unique<LimitNode>();
+ limitNode->children.push_back(virtScan.release());
+ limitNode->limit = 1;
+
+ // Make a QuerySolution from the root limitNode.
+ auto querySolution = makeQuerySolution(std::move(limitNode));
+
+ // Translate the QuerySolution tree to an sbe::PlanStage.
+ auto [resultSlots, stage, data] = buildPlanStage(std::move(querySolution), true);
+
+ // Prepare the sbe::PlanStage for execution.
+ auto resultAccessors = prepareTree(&data.ctx, stage.get(), resultSlots);
+
+ int64_t index = 0;
+ for (auto st = stage->getNext(); st == sbe::PlanState::ADVANCED; st = stage->getNext()) {
+ // Assert that the recordIDs are what we expect.
+ auto [tag, val] = resultAccessors[0]->getViewOfValue();
+ ASSERT_TRUE(tag == sbe::value::TypeTags::NumberInt64);
+ ASSERT_EQ(index, sbe::value::bitcastTo<int64_t>(val));
+
+ // Assert that the document produced from the stage is what we expect.
+ auto [tagDoc, valDoc] = resultAccessors[1]->getViewOfValue();
+ ASSERT_TRUE(tagDoc == sbe::value::TypeTags::bsonObject);
+ auto bo = BSONObj(sbe::value::bitcastTo<const char*>(valDoc));
+
+ ASSERT_BSONOBJ_EQ(bo, BSON("a" << ++index << "b" << 2));
+ }
+ ASSERT_EQ(index, 1);
+}
+} // namespace mongo
diff --git a/src/mongo/db/query/sbe_stage_builder_test_fixture.cpp b/src/mongo/db/query/sbe_stage_builder_test_fixture.cpp
new file mode 100644
index 00000000000..8762c3c0f5d
--- /dev/null
+++ b/src/mongo/db/query/sbe_stage_builder_test_fixture.cpp
@@ -0,0 +1,74 @@
+/**
+ * Copyright (C) 2020-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/catalog/collection.h"
+#include "mongo/db/matcher/extensions_callback_noop.h"
+#include "mongo/db/pipeline/expression_context_for_test.h"
+#include "mongo/db/query/mock_yield_policies.h"
+#include "mongo/db/query/sbe_stage_builder.h"
+#include "mongo/db/query/sbe_stage_builder_test_fixture.h"
+
+namespace mongo {
+std::unique_ptr<QuerySolution> SBEStageBuilderTestFixture::makeQuerySolution(
+ std::unique_ptr<QuerySolutionNode> root) {
+ auto querySoln = std::make_unique<QuerySolution>();
+ querySoln->setRoot(std::move(root));
+ return querySoln;
+}
+
+std::tuple<sbe::value::SlotVector, std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData>
+SBEStageBuilderTestFixture::buildPlanStage(std::unique_ptr<QuerySolution> querySolution,
+ bool hasRecordId) {
+ auto qr = std::make_unique<QueryRequest>(_nss);
+ const boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest(_nss));
+ auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr), expCtx);
+ ASSERT_OK(statusWithCQ.getStatus());
+
+ stage_builder::SlotBasedStageBuilder builder{opCtx(),
+ CollectionPtr::null,
+ *statusWithCQ.getValue(),
+ *querySolution,
+ nullptr /* YieldPolicy */,
+ false};
+
+ auto stage = builder.build(querySolution->root());
+ auto data = builder.getPlanStageData();
+
+ auto slots = sbe::makeSV();
+ if (hasRecordId) {
+ slots.push_back(*data.recordIdSlot);
+ }
+ slots.push_back(*data.resultSlot);
+
+ return {slots, std::move(stage), std::move(data)};
+}
+
+} // namespace mongo
diff --git a/src/mongo/db/query/sbe_stage_builder_test_fixture.h b/src/mongo/db/query/sbe_stage_builder_test_fixture.h
new file mode 100644
index 00000000000..d38382e8c71
--- /dev/null
+++ b/src/mongo/db/query/sbe_stage_builder_test_fixture.h
@@ -0,0 +1,79 @@
+/**
+ * Copyright (C) 2020-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.
+ */
+
+#pragma once
+
+#include "mongo/db/exec/sbe/sbe_plan_stage_test.h"
+#include "mongo/db/exec/sbe/stages/stages.h"
+#include "mongo/db/exec/sbe/values/slot.h"
+#include "mongo/db/exec/sbe/values/value.h"
+#include "mongo/db/query/query_solution.h"
+#include "mongo/unittest/unittest.h"
+
+namespace mongo {
+/**
+ * SBEStageBuilderTestFixture is a unittest fixture that can be used to facilitate testing the
+ * translation of a QuerySolution tree to an sbe PlanStage tree.
+ *
+ * The main mechanism that enables the whole sbe::PlanStage tree to be exercised under unittests is
+ * the use of a VirtualScanNode. This virtual node can be created with a vector of BSONArray
+ * documents and used as a replacement for a CollectionScanNode in the QuerySolution tree. A testing
+ * client would manually build a QuerySolution tree containing this VirtualScanNode and
+ * then transform it to an sbe::PlanStage by calling buildPlanStage(). The buildPlanStage()
+ * method will do the QuerySolution to sbe::PlanStage translation, and return a vector of result
+ * slots, the prepared sub-tree and an sbe::PlanStageData that carries the sbe::CompileCtx needed to
+ * prepare the sbe::PlanStage tree. The sbe::PlanStageData returned from buildPlanStage() must be
+ * kept alive across buildPlanStage(), prepareTree() and execution of the plan.
+ */
+class SBEStageBuilderTestFixture : public sbe::PlanStageTestFixture {
+public:
+ SBEStageBuilderTestFixture() = default;
+
+ /**
+ * Makes a QuerySolution from a QuerySolutionNode.
+ */
+ std::unique_ptr<QuerySolution> makeQuerySolution(std::unique_ptr<QuerySolutionNode> root);
+
+ /**
+ * Builds an sbe::PlanStage tree from a QuerySolution that can be executed by repeatedly calling
+ * getNext() on the root. Results from the PlanStage are exposed through the returned
+ * SlotVector. If hasRecordId is 'true' then the returned SlotVector will carry a SlotId in the
+ * 0th position for the RecordId and a SlotId for the BSONObj representation of the document in
+ * the 1st position. Otherwise, if hasRecordId is 'false', the SlotVector will contain a single
+ * SlotId for the BSONObj representation of the document.
+ */
+ std::
+ tuple<sbe::value::SlotVector, std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData>
+ buildPlanStage(std::unique_ptr<QuerySolution> querySolution, bool hasRecordId);
+
+private:
+ const NamespaceString _nss = NamespaceString{"testdb.sbe_stage_builder"};
+};
+
+} // namespace mongo
diff --git a/src/mongo/db/query/stage_types.h b/src/mongo/db/query/stage_types.h
index 6b16e4575d0..882d0bf032b 100644
--- a/src/mongo/db/query/stage_types.h
+++ b/src/mongo/db/query/stage_types.h
@@ -52,6 +52,10 @@ enum StageType {
STAGE_CACHED_PLAN,
STAGE_COLLSCAN,
+ // A virtual scan stage that simulates a collection scan and doesn't depend on underlying
+ // storage.
+ STAGE_VIRTUAL_SCAN,
+
// This stage sits at the root of the query tree and counts up the number of results
// returned by its child.
STAGE_COUNT,