diff options
Diffstat (limited to 'src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp')
-rw-r--r-- | src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp | 239 |
1 files changed, 239 insertions, 0 deletions
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp new file mode 100644 index 00000000000..f9a70523ee0 --- /dev/null +++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp @@ -0,0 +1,239 @@ +/** + * Copyright (C) 2018-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. + */ + +/** + * This file contains a test framework for testing sbe::PlanStages. + */ + +#include "mongo/platform/basic.h" + +#include "mongo/db/exec/sbe/sbe_plan_stage_test.h" + +#include <string_view> + +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(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(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>(), 1, boost::none), + 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. + + // 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(int64_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 (int64_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::NumberInt64, i)))); + } + + return {std::move(projectSlots), + makeS<ProjectStage>(std::move(scanStage), std::move(projections))}; +} + +std::pair<value::SlotId, std::unique_ptr<PlanStage>> PlanStageTestFixture::generateMockScan( + const BSONArray& array) { + auto [arrTag, arrVal] = makeValue(array); + return generateMockScan(arrTag, arrVal); +} + +std::pair<value::SlotVector, std::unique_ptr<PlanStage>> +PlanStageTestFixture::generateMockScanMulti(int64_t numSlots, const BSONArray& array) { + auto [arrTag, arrVal] = makeValue(array); + return generateMockScanMulti(numSlots, arrTag, arrVal); +} + +void PlanStageTestFixture::prepareTree(PlanStage* root) { + root->prepare(*_compileCtx); + root->attachFromOperationContext(opCtx()); + root->open(false); +} + +value::SlotAccessor* PlanStageTestFixture::prepareTree(PlanStage* root, value::SlotId slot) { + prepareTree(root); + return root->getAccessor(*_compileCtx, slot); +} + +std::vector<value::SlotAccessor*> PlanStageTestFixture::prepareTree(PlanStage* root, + value::SlotVector slots) { + std::vector<value::SlotAccessor*> slotAccessors; + + prepareTree(root); + for (auto slot : slots) { + slotAccessors.emplace_back(root->getAccessor(*_compileCtx, slot)); + } + return slotAccessors; +} + +std::pair<value::TypeTags, value::Value> PlanStageTestFixture::getAllResults( + PlanStage* stage, value::SlotAccessor* accessor) { + // Allocate an array to hold the results. + auto [resultsTag, resultsVal] = value::makeNewArray(); + value::ValueGuard guard{resultsTag, resultsVal}; + auto resultsView = value::getArrayView(resultsVal); + // Loop and repeatedly call getNext() until we reach the end, storing the values produced + // into the array. + for (auto st = stage->getNext(); st == PlanState::ADVANCED; st = stage->getNext()) { + auto [tag, val] = accessor->copyOrMoveValue(); + resultsView->push_back(tag, val); + } + + guard.reset(); + return {resultsTag, resultsVal}; +} + +std::pair<value::TypeTags, value::Value> PlanStageTestFixture::getAllResultsMulti( + PlanStage* stage, std::vector<value::SlotAccessor*> accessors) { + // Allocate an SBE array to hold the results. + auto [resultsTag, resultsVal] = value::makeNewArray(); + value::ValueGuard resultsGuard{resultsTag, resultsVal}; + auto resultsView = value::getArrayView(resultsVal); + + // Loop and repeatedly call getNext() until we reach the end. + for (auto st = stage->getNext(); st == PlanState::ADVANCED; st = stage->getNext()) { + // Create a new SBE array (`arr`) containing the values produced by each SlotAccessor + // and insert `arr` into the array of results. + auto [arrTag, arrVal] = value::makeNewArray(); + value::ValueGuard guard{arrTag, arrVal}; + auto arrView = value::getArrayView(arrVal); + for (size_t i = 0; i < accessors.size(); ++i) { + auto [tag, val] = accessors[i]->copyOrMoveValue(); + arrView->push_back(tag, val); + } + guard.reset(); + resultsView->push_back(arrTag, arrVal); + } + + resultsGuard.reset(); + return {resultsTag, resultsVal}; +} + +void PlanStageTestFixture::runTest(value::TypeTags inputTag, + value::Value inputVal, + value::TypeTags expectedTag, + value::Value expectedVal, + const MakeStageFn<value::SlotId>& makeStage) { + // 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); + + // 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); + + // Get all the results produced by the PlanStage we want to test. + auto [resultsTag, resultsVal] = getAllResults(stage.get(), resultAccessor); + value::ValueGuard resultGuard{resultsTag, resultsVal}; + + // Compare the results produced with the expected output and assert that they match. + ASSERT_TRUE(valueEquals(resultsTag, resultsVal, expectedTag, expectedVal)); +} + +void PlanStageTestFixture::runTestMulti(int64_t numInputSlots, + value::TypeTags inputTag, + value::Value inputVal, + value::TypeTags expectedTag, + value::Value expectedVal, + const MakeStageFn<value::SlotVector>& makeStageMulti) { + // 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); + + // 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); + + // Get all the results produced by the PlanStage we want to test. + auto [resultsTag, resultsVal] = getAllResultsMulti(stage.get(), resultAccessors); + value::ValueGuard resultGuard{resultsTag, resultsVal}; + + // Compare the results produced with the expected output and assert that they match. + ASSERT_TRUE(valueEquals(resultsTag, resultsVal, expectedTag, expectedVal)); +} + +} // namespace mongo::sbe |