diff options
Diffstat (limited to 'src')
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, |