diff options
author | Davis Haupt <davis.haupt@mongodb.com> | 2023-01-24 21:48:44 +0000 |
---|---|---|
committer | Evergreen Agent <no-reply@evergreen.mongodb.com> | 2023-01-24 23:06:29 +0000 |
commit | 07d55f78e9cbb26680f4ed3e1651648faefced50 (patch) | |
tree | 2f336548f19823705a1a4897c267f9b65f82d2d7 /src | |
parent | 432d6f8687c8b6dc6471241d1a7fb1cc04970899 (diff) | |
download | mongo-07d55f78e9cbb26680f4ed3e1651648faefced50.tar.gz |
SERVER-73240 readability improvements in abt_lower
Diffstat (limited to 'src')
-rw-r--r-- | src/mongo/db/exec/sbe/abt/abt_lower.cpp | 183 | ||||
-rw-r--r-- | src/mongo/db/exec/sbe/abt/abt_lower.h | 21 |
2 files changed, 111 insertions, 93 deletions
diff --git a/src/mongo/db/exec/sbe/abt/abt_lower.cpp b/src/mongo/db/exec/sbe/abt/abt_lower.cpp index 0c047825e55..bd6b5d0fd3a 100644 --- a/src/mongo/db/exec/sbe/abt/abt_lower.cpp +++ b/src/mongo/db/exec/sbe/abt/abt_lower.cpp @@ -76,7 +76,48 @@ std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport(const Constan } std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport(const Source&) { - uasserted(6624202, "not yet implemented"); + tasserted(6624202, "not yet implemented"); + return nullptr; +} + +void SBEExpressionLowering::prepare(const Let& let) { + // Assign a frame ID for the local variable bound by this Let expression. + _letMap[&let] = ++_frameCounter; +} + +std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( + const Let& let, std::unique_ptr<sbe::EExpression> bind, std::unique_ptr<sbe::EExpression> in) { + auto it = _letMap.find(&let); + tassert(6624206, "incorrect let map", it != _letMap.end()); + auto frameId = it->second; + _letMap.erase(it); + + // ABT let binds only a single variable. When we extend it to support multiple binds then we + // have to revisit how we map variable names to sbe slot ids. + return sbe::makeE<sbe::ELocalBind>(frameId, sbe::makeEs(std::move(bind)), std::move(in)); +} + +void SBEExpressionLowering::prepare(const LambdaAbstraction& lam) { + // Assign a frame ID for the local variable bound by this LambdaAbstraction. + _lambdaMap[&lam] = ++_frameCounter; +} + +std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( + const LambdaAbstraction& lam, std::unique_ptr<sbe::EExpression> body) { + auto it = _lambdaMap.find(&lam); + tassert(6624207, "incorrect lambda map", it != _lambdaMap.end()); + auto frameId = it->second; + _lambdaMap.erase(it); + + return sbe::makeE<sbe::ELocalLambda>(frameId, std::move(body)); +} + +std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( + const LambdaApplication&, + std::unique_ptr<sbe::EExpression> lam, + std::unique_ptr<sbe::EExpression> arg) { + // lambda applications are not directly supported by SBE (yet) and must not be present. + tasserted(6624208, "lambda application is not implemented"); return nullptr; } @@ -84,24 +125,29 @@ std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport(const Variabl auto def = _env.getDefinition(var); if (!def.definedBy.empty()) { + // If this variable was defined by a Let expression, use the frame ID that was defined in + // the prepare() step for the Let. if (auto let = def.definedBy.cast<Let>(); let) { auto it = _letMap.find(let); - uassert(6624203, "incorrect let map", it != _letMap.end()); + tassert(6624203, "incorrect let map", it != _letMap.end()); return sbe::makeE<sbe::EVariable>(it->second, 0, _env.isLastRef(var)); } else if (auto lam = def.definedBy.cast<LambdaAbstraction>(); lam) { - // This is a lambda parameter. + // Similarly if the variable was defined by a lambda abstraction, use a frame ID rather + // than a slot. auto it = _lambdaMap.find(lam); - uassert(6624204, "incorrect lambda map", it != _lambdaMap.end()); + tassert(6624204, "incorrect lambda map", it != _lambdaMap.end()); return sbe::makeE<sbe::EVariable>(it->second, 0, _env.isLastRef(var)); } } + // If variable was not defined in the scope of the local expression via a Let or + // LambdaAbstraction, it must be a reference that will be in the slotMap. if (auto it = _slotMap.find(var.name()); it != _slotMap.end()) { // Found the slot. return sbe::makeE<sbe::EVariable>(it->second); } - uasserted(6624205, str::stream() << "undefined variable: " << var.name()); + tasserted(6624205, str::stream() << "undefined variable: " << var.name()); return nullptr; } @@ -182,45 +228,6 @@ std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( return sbe::makeE<sbe::EIf>(std::move(cond), std::move(thenBranch), std::move(elseBranch)); } -void SBEExpressionLowering::prepare(const Let& let) { - _letMap[&let] = ++_frameCounter; -} - -std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( - const Let& let, std::unique_ptr<sbe::EExpression> bind, std::unique_ptr<sbe::EExpression> in) { - auto it = _letMap.find(&let); - uassert(6624206, "incorrect let map", it != _letMap.end()); - auto frameId = it->second; - _letMap.erase(it); - - // ABT let binds only a single variable. When we extend it to support multiple binds then we - // have to revisit how we map variable names to sbe slot ids. - return sbe::makeE<sbe::ELocalBind>(frameId, sbe::makeEs(std::move(bind)), std::move(in)); -} - -void SBEExpressionLowering::prepare(const LambdaAbstraction& lam) { - _lambdaMap[&lam] = ++_frameCounter; -} - -std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( - const LambdaAbstraction& lam, std::unique_ptr<sbe::EExpression> body) { - auto it = _lambdaMap.find(&lam); - uassert(6624207, "incorrect lambda map", it != _lambdaMap.end()); - auto frameId = it->second; - _lambdaMap.erase(it); - - return sbe::makeE<sbe::ELocalLambda>(frameId, std::move(body)); -} - -std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( - const LambdaApplication&, - std::unique_ptr<sbe::EExpression> lam, - std::unique_ptr<sbe::EExpression> arg) { - // lambda applications are not directly supported by SBE (yet) and must not be present. - uasserted(6624208, "lambda application is not implemented"); - return nullptr; -} - std::unique_ptr<sbe::EExpression> SBEExpressionLowering::transport( const FunctionCall& fn, std::vector<std::unique_ptr<sbe::EExpression>> args) { auto name = fn.name(); @@ -301,7 +308,7 @@ sbe::value::SlotVector SBENodeLowering::convertProjectionsToSlots( sbe::value::SlotVector result; for (const ProjectionName& projectionName : projectionNames) { auto it = _slotMap.find(projectionName); - uassert(6624211, + tassert(6624211, str::stream() << "undefined variable: " << projectionName, it != _slotMap.end()); result.push_back(it->second); @@ -355,12 +362,12 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const RootNode& n, auto input = generateInternal(child); auto output = refs.cast<References>(); - uassert(6624212, "refs expected", output); + tassert(6624212, "refs expected", output); SlotVarMap finalMap; for (auto& o : output->nodes()) { auto var = o.cast<Variable>(); - uassert(6624213, "var expected", var); + tassert(6624213, "var expected", var); if (auto it = _slotMap.find(var->name()); it != _slotMap.end()) { finalMap.emplace(var->name(), it->second); } @@ -389,22 +396,20 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const EvaluationNode& n, const ABT& binds) { auto input = generateInternal(child); + // If the evaluation node is only renaming a variable, do not place a project stage. if (auto varPtr = n.getProjection().cast<Variable>(); varPtr != nullptr) { - // Evaluation node is only renaming a variable. Do not place a project stage. mapProjToSlot(n.getProjectionName(), _slotMap.at(varPtr->name())); return input; } - auto binder = binds.cast<ExpressionBinder>(); - uassert(6624214, "binder expected", binder); - - auto& names = binder->names(); - auto& exprs = binder->exprs(); + auto& binder = n.binder(); + auto& names = binder.names(); + auto& exprs = binder.exprs(); sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projects; for (size_t idx = 0; idx < exprs.size(); ++idx) { - auto expr = SBEExpressionLowering{_env, _slotMap, _namedSlots}.optimize(exprs[idx]); + auto expr = lowerExpression(exprs[idx]); auto slot = _slotIdGenerator.generate(); mapProjToSlot(names[idx], slot); @@ -419,7 +424,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const FilterNode& n, const ABT& child, const ABT& filter) { auto input = generateInternal(child); - auto expr = SBEExpressionLowering{_env, _slotMap, _namedSlots}.optimize(filter); + auto expr = lowerExpression(filter); const PlanNodeId planNodeId = _nodeToGroupPropsMap.at(&n)._planNodeId; // Check if the filter expression is 'constant' (i.e., does not depend on any variables); then @@ -451,7 +456,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const ExchangeNode& n, // The DOP is obtained from the child (number of producers). const auto& childProps = _nodeToGroupPropsMap.at(n.getChild().cast<Node>())._physicalProps; const auto& childDistribution = getPropertyConst<DistributionRequirement>(childProps); - uassert(6624330, + tassert(6624330, "Parent and child distributions are the same", !(childDistribution == n.getProperty())); @@ -459,7 +464,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const ExchangeNode& n, (childDistribution.getDistributionAndProjections()._type == DistributionType::Centralized) ? 1 : _metadata._numberOfPartitions; - uassert(6624215, "invalid DOP", localDOP >= 1); + tassert(6624215, "invalid DOP", localDOP >= 1); auto input = generateInternal(child); @@ -491,7 +496,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const ExchangeNode& n, std::vector<std::unique_ptr<sbe::EExpression>> args; for (const ProjectionName& proj : distribAndProjections._projectionNames) { auto it = _slotMap.find(proj); - uassert(6624216, str::stream() << "undefined var: " << proj, it != _slotMap.end()); + tassert(6624216, str::stream() << "undefined var: " << proj, it != _slotMap.end()); args.emplace_back(sbe::makeE<sbe::EVariable>(it->second)); } @@ -500,7 +505,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const ExchangeNode& n, } case DistributionType::UnknownPartitioning: - uasserted(6624217, "Cannot partition into unknown distribution"); + tasserted(6624217, "Cannot partition into unknown distribution"); default: MONGO_UNREACHABLE; @@ -520,6 +525,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const ExchangeNode& n, static sbe::value::SortDirection collationOpToSBESortDirection(const CollationOp collOp) { switch (collOp) { + // TODO: is there a more efficient way to compute clustered collation op than sort? case CollationOp::Ascending: case CollationOp::Clustered: return sbe::value::SortDirection::Ascending; @@ -536,17 +542,14 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const CollationNode& n, sbe::value::SlotVector orderBySlots; std::vector<sbe::value::SortDirection> directions; - ProjectionNameVector collationProjections; for (const auto& entry : n.getProperty().getCollationSpec()) { - collationProjections.push_back(entry.first); auto it = _slotMap.find(entry.first); - uassert(6624219, + tassert(6624219, str::stream() << "undefined orderBy var: " << entry.first, it != _slotMap.end()); orderBySlots.push_back(it->second); - // TODO: is there a more efficient way to compute clustered collation op than sort? directions.push_back(collationOpToSBESortDirection(entry.second)); } @@ -557,7 +560,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const CollationNode& n, if (properties::hasProperty<properties::LimitSkipRequirement>(physProps)) { const auto& limitSkipReq = properties::getPropertyConst<properties::LimitSkipRequirement>(physProps); - uassert(6624221, "We should not have skip set here", limitSkipReq.getSkip() == 0); + tassert(6624221, "We should not have skip set here", limitSkipReq.getSkip() == 0); limit = limitSkipReq.getLimit(); } @@ -584,7 +587,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const UniqueNode& n, sbe::value::SlotVector keySlots; for (const ProjectionName& projectionName : n.getProjections()) { auto it = _slotMap.find(projectionName); - uassert(6624222, + tassert(6624222, str::stream() << "undefined variable: " << projectionName, it != _slotMap.end()); keySlots.push_back(it->second); @@ -604,7 +607,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const SpoolProducerNode& n sbe::value::SlotVector vals; for (const ProjectionName& projectionName : n.binder().names()) { auto it = _slotMap.find(projectionName); - uassert(6624139, + tassert(6624139, str::stream() << "undefined variable: " << projectionName, it != _slotMap.end()); vals.push_back(it->second); @@ -617,7 +620,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const SpoolProducerNode& n std::move(input), n.getSpoolId(), std::move(vals), planNodeId); case SpoolProducerType::Lazy: { - auto expr = SBEExpressionLowering{_env, _slotMap, _namedSlots}.optimize(filter); + auto expr = lowerExpression(filter); return sbe::makeS<sbe::SpoolLazyProducerStage>( std::move(input), n.getSpoolId(), std::move(vals), std::move(expr), planNodeId); } @@ -664,28 +667,26 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const GroupByNode& n, // represent both so that distinction is kind of moot. sbe::value::SlotVector gbs; auto gbCols = gbRefs.cast<References>(); - uassert(6624223, "refs expected", gbCols); + tassert(6624223, "refs expected", gbCols); for (auto& o : gbCols->nodes()) { auto var = o.cast<Variable>(); - uassert(6624224, "var expected", var); + tassert(6624224, "var expected", var); auto it = _slotMap.find(var->name()); - uassert(6624225, str::stream() << "undefined var: " << var->name(), it != _slotMap.end()); + tassert(6624225, str::stream() << "undefined var: " << var->name(), it != _slotMap.end()); gbs.push_back(it->second); } // Similar considerations apply to the agg expressions as to the group by columns. - auto binderAgg = aggBinds.cast<ExpressionBinder>(); - uassert(6624226, "binder expected", binderAgg); - auto refsAgg = aggRefs.cast<References>(); - uassert(6624227, "refs expected", refsAgg); + auto& names = n.binderAgg().names(); - auto& names = binderAgg->names(); + auto refsAgg = aggRefs.cast<References>(); + tassert(6624227, "refs expected", refsAgg); auto& exprs = refsAgg->nodes(); sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> aggs; for (size_t idx = 0; idx < exprs.size(); ++idx) { - auto expr = SBEExpressionLowering{_env, _slotMap, _namedSlots}.optimize(exprs[idx]); + auto expr = lowerExpression(exprs[idx]); auto slot = _slotIdGenerator.generate(); mapProjToSlot(names[idx], slot); @@ -720,7 +721,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const NestedLoopJoinNode& correlatedSlots.push_back(_slotMap.at(projectionName)); } - auto expr = SBEExpressionLowering{_env, _slotMap, _namedSlots}.optimize(filter); + auto expr = lowerExpression(filter); const auto& leftChildProps = _nodeToGroupPropsMap.at(n.getLeftChild().cast<Node>()); auto outerProjects = convertRequiredProjectionsToSlots(leftChildProps); @@ -760,7 +761,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const HashJoinNode& n, auto innerStage = generateInternal(leftChild); auto outerStage = generateInternal(rightChild); - uassert(6624228, "Only inner joins supported for now", n.getJoinType() == JoinType::Inner); + tassert(6624228, "Only inner joins supported for now", n.getJoinType() == JoinType::Inner); const auto& leftProps = _nodeToGroupPropsMap.at(n.getLeftChild().cast<Node>()); const auto& rightProps = _nodeToGroupPropsMap.at(n.getRightChild().cast<Node>()); @@ -820,9 +821,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const SortedMergeNode& n, const ABTVector& children, const ABT& binder, const ABT& /*refs*/) { - const auto exprBinder = binder.cast<ExpressionBinder>(); - uassert(7063705, "binder expected", exprBinder); - const auto& names = exprBinder->names(); + const auto& names = n.binder().names(); const ProjectionCollationSpec& collSpec = n.getCollationReq().getCollationSpec(); @@ -892,9 +891,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const UnionNode& n, const ABTVector& children, const ABT& binder, const ABT& /*refs*/) { - auto unionBinder = binder.cast<ExpressionBinder>(); - uassert(6624229, "binder expected", unionBinder); - const auto& names = unionBinder->names(); + const auto& names = n.binder().names(); sbe::PlanStage::Vector loweredChildren; std::vector<sbe::value::SlotVector> inputVals; @@ -950,7 +947,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const UnwindNode& n, auto input = generateInternal(child); auto it = _slotMap.find(n.getProjectionName()); - uassert(6624230, + tassert(6624230, str::stream() << "undefined unwind variable: " << n.getProjectionName(), it != _slotMap.end()); @@ -999,7 +996,7 @@ static NamespaceStringOrUUID parseFromScanDef(const ScanDefinition& def) { std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const PhysicalScanNode& n, const ABT& /*binds*/) { const ScanDefinition& def = _metadata._scanDefs.at(n.getScanDefName()); - uassert(6624231, "Collection must exist to lower Scan", def.exists()); + tassert(6624231, "Collection must exist to lower Scan", def.exists()); auto& typeSpec = def.getOptionsMap().at("type"); boost::optional<sbe::value::SlotId> ridSlot; @@ -1058,7 +1055,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const PhysicalScanNode& n, callbacks, _scanOrder == ScanOrder::Random); } else { - uasserted(6624355, "Unknown scan type."); + tasserted(6624355, "Unknown scan type."); } return nullptr; } @@ -1082,7 +1079,7 @@ std::unique_ptr<sbe::EExpression> SBENodeLowering::convertBoundsToExpr( sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt32, sbe::value::bitcastFrom<uint32_t>(indexDef.getOrdering()))); - auto exprLower = SBEExpressionLowering{_env, _slotMap, _namedSlots}; + auto exprLower = getExpressionLowering(); bool inclusive = true; bool fullyInfinite = true; for (const auto& entry : interval) { @@ -1114,7 +1111,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const IndexScanNode& n, co const std::string& indexDefName = n.getIndexDefName(); const ScanDefinition& scanDef = _metadata._scanDefs.at(n.getScanDefName()); - uassert(6624232, "Collection must exist to lower IndexScan", scanDef.exists()); + tassert(6624232, "Collection must exist to lower IndexScan", scanDef.exists()); const IndexDefinition& indexDef = scanDef.getIndexDefs().at(indexDefName); NamespaceStringOrUUID nss = parseFromScanDef(scanDef); @@ -1124,7 +1121,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const IndexScanNode& n, co std::vector<std::string> fields; sbe::value::SlotVector vars; generateSlots(fieldProjectionMap, ridSlot, rootSlot, fields, vars); - uassert(6624233, "Cannot deliver root projection in this context", !rootSlot.has_value()); + tassert(6624233, "Cannot deliver root projection in this context", !rootSlot.has_value()); std::vector<std::pair<size_t, sbe::value::SlotId>> indexVars; sbe::IndexKeysInclusionSet indexKeysToInclude; @@ -1145,7 +1142,7 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const IndexScanNode& n, co const auto& interval = n.getIndexInterval(); auto lowerBoundExpr = convertBoundsToExpr(true /*isLower*/, indexDef, interval); auto upperBoundExpr = convertBoundsToExpr(false /*isLower*/, indexDef, interval); - uassert(6624234, + tassert(6624234, "Invalid bounds combination", lowerBoundExpr != nullptr || upperBoundExpr == nullptr); @@ -1178,10 +1175,10 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const SeekNode& n, const ABT& /*binds*/, const ABT& /*refs*/) { const ScanDefinition& def = _metadata._scanDefs.at(n.getScanDefName()); - uassert(6624235, "Collection must exist to lower Seek", def.exists()); + tassert(6624235, "Collection must exist to lower Seek", def.exists()); auto& typeSpec = def.getOptionsMap().at("type"); - uassert(6624236, "SeekNode only supports mongod collections", typeSpec == "mongod"); + tassert(6624236, "SeekNode only supports mongod collections", typeSpec == "mongod"); NamespaceStringOrUUID nss = parseFromScanDef(def); boost::optional<sbe::value::SlotId> ridSlot; diff --git a/src/mongo/db/exec/sbe/abt/abt_lower.h b/src/mongo/db/exec/sbe/abt/abt_lower.h index 4943144938f..b85c8bf43b5 100644 --- a/src/mongo/db/exec/sbe/abt/abt_lower.h +++ b/src/mongo/db/exec/sbe/abt/abt_lower.h @@ -202,7 +202,18 @@ private: std::vector<std::string>& fields, sbe::value::SlotVector& vars); + /** + * Convert a vector of ProjectionNames to slot IDs from the projections that have already been + * bound to slots. + */ sbe::value::SlotVector convertProjectionsToSlots(const ProjectionNameVector& projectionNames); + + /** + * During Cascades, projections that a node is required to propagate up the tree are added to + * the RequiredProjections node property. This function pulls out those projection names and + * looks up the relevant slot IDs they are bound to. The optional toExclude vector can prevent + * some slots from being added to the output vector. + */ sbe::value::SlotVector convertRequiredProjectionsToSlots( const NodeProps& props, const sbe::value::SlotVector& toExclude = {}); @@ -224,6 +235,16 @@ private: sbe::value::SlotId slot, bool canOverwrite = false); + /** + * Instantiate an expression lowering transporter for use in node lowering. + */ + SBEExpressionLowering getExpressionLowering() { + return SBEExpressionLowering{_env, _slotMap, _namedSlots}; + } + + std::unique_ptr<sbe::EExpression> lowerExpression(const ABT& e) { + return getExpressionLowering().optimize(e); + } const VariableEnvironment& _env; SlotVarMap& _slotMap; const NamedSlotsProvider& _namedSlots; |