diff options
Diffstat (limited to 'src/mongo/db/query/sbe_stage_builder_helpers.h')
-rw-r--r-- | src/mongo/db/query/sbe_stage_builder_helpers.h | 87 |
1 files changed, 74 insertions, 13 deletions
diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.h b/src/mongo/db/query/sbe_stage_builder_helpers.h index f3b5d9ca455..634005c504c 100644 --- a/src/mongo/db/query/sbe_stage_builder_helpers.h +++ b/src/mongo/db/query/sbe_stage_builder_helpers.h @@ -44,6 +44,10 @@ #include "mongo/db/query/sbe_stage_builder_eval_frame.h" #include "mongo/db/query/stage_types.h" +namespace mongo::projection_ast { +class Projection; +} + namespace mongo::stage_builder { std::unique_ptr<sbe::EExpression> makeUnaryOp(sbe::EPrimUnary::Op unaryOp, @@ -524,19 +528,20 @@ std::unique_ptr<sbe::EExpression> makeLocalBind(sbe::value::FrameIdGenerator* fr return sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(innerExpr)); } - -std::tuple<sbe::value::SlotId, sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> -makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage, - sbe::value::SlotId seekKeySlot, - sbe::value::SlotId snapshotIdSlot, - sbe::value::SlotId indexIdSlot, - sbe::value::SlotId indexKeySlot, - sbe::value::SlotId indexKeyPatternSlot, - const CollectionPtr& collToFetch, - StringMap<const IndexAccessMethod*> iamMap, - PlanNodeId planNodeId, - sbe::value::SlotVector slotsToForward, - sbe::value::SlotIdGenerator& slotIdGenerator); +std::unique_ptr<sbe::PlanStage> makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inputStage, + sbe::value::SlotId resultSlot, + sbe::value::SlotId recordIdSlot, + std::vector<std::string> fields, + sbe::value::SlotVector fieldSlots, + sbe::value::SlotId seekKeySlot, + sbe::value::SlotId snapshotIdSlot, + sbe::value::SlotId indexIdSlot, + sbe::value::SlotId indexKeySlot, + sbe::value::SlotId indexKeyPatternSlot, + const CollectionPtr& collToFetch, + StringMap<const IndexAccessMethod*> iamMap, + PlanNodeId planNodeId, + sbe::value::SlotVector slotsToForward); /** * Trees generated by 'generateFilter' maintain state during execution. There are two types of state @@ -1010,4 +1015,60 @@ std::unique_ptr<sbe::PlanStage> rehydrateIndexKey(std::unique_ptr<sbe::PlanStage const sbe::value::SlotVector& indexKeySlots, sbe::value::SlotId resultSlot); +IndexKeyPatternTreeNode buildPatternTree(const projection_ast::Projection& projection); + +/** + * This method retrieves the values of the specified top-level fields ('fields') from 'resultSlot' + * and stores the values into slots. + * + * This method returns a pair containing: (1) the updated SBE plan stage tree and; (2) a vector of + * the slots ('slots') containing the field values. + * + * The order of slots in 'slots' will match the order of fields in 'fields'. + */ +std::pair<std::unique_ptr<sbe::PlanStage>, sbe::value::SlotVector> projectTopLevelFields( + std::unique_ptr<sbe::PlanStage> stage, + const std::vector<std::string>& fields, + sbe::value::SlotId resultSlot, + PlanNodeId nodeId, + sbe::value::SlotIdGenerator* slotIdGenerator); + +/** + * This method projects the constant value Nothing to multiple slots (the specific number of slots + * being specified by parameter 'n'). + * + * This method returns a pair containing: (1) the updated SBE plan stage tree and; (2) a vector of + * slots ('slots') containing Nothing. + * + * The number of slots in 'slots' will always be equal to parameter 'n'. + */ +std::pair<std::unique_ptr<sbe::PlanStage>, sbe::value::SlotVector> projectNothingToSlots( + std::unique_ptr<sbe::PlanStage> stage, + size_t n, + PlanNodeId nodeId, + sbe::value::SlotIdGenerator* slotIdGenerator); + +template <typename T, typename FuncT> +std::vector<T> filterVector(std::vector<T> vec, FuncT fn) { + std::vector<T> result; + std::copy_if(std::make_move_iterator(vec.begin()), + std::make_move_iterator(vec.end()), + std::back_inserter(result), + fn); + return result; +} + +template <typename T, typename FuncT> +std::pair<std::vector<T>, std::vector<T>> splitVector(std::vector<T> vec, FuncT fn) { + std::pair<std::vector<T>, std::vector<T>> result; + for (size_t i = 0; i < vec.size(); ++i) { + if (fn(vec[i])) { + result.first.emplace_back(std::move(vec[i])); + } else { + result.second.emplace_back(std::move(vec[i])); + } + } + return result; +} + } // namespace mongo::stage_builder |