diff options
author | Eric Cox <eric.cox@mongodb.com> | 2022-06-24 13:49:26 +0000 |
---|---|---|
committer | Eric Cox <eric.cox@mongodb.com> | 2022-06-24 13:49:26 +0000 |
commit | e27fb371450c1aecbf3045c13c9a5257560ee615 (patch) | |
tree | 7813bd1533e7d45af38ea91879468f74974ac2fb | |
parent | d8ae3f7be6893c89be0ac15a0aa2f1073d9e9cb2 (diff) | |
download | mongo-e27fb371450c1aecbf3045c13c9a5257560ee615.tar.gz |
wip refactor stage builder
-rw-r--r-- | src/mongo/db/exec/sbe/abt/abt_lower.cpp | 47 | ||||
-rw-r--r-- | src/mongo/db/exec/sbe/stages/ix_scan.cpp | 111 | ||||
-rw-r--r-- | src/mongo/db/exec/sbe/stages/ix_scan.h | 38 | ||||
-rw-r--r-- | src/mongo/db/query/sbe_stage_builder_index_scan.cpp | 46 | ||||
-rw-r--r-- | src/mongo/db/query/sbe_stage_builder_lookup.cpp | 4 |
5 files changed, 83 insertions, 163 deletions
diff --git a/src/mongo/db/exec/sbe/abt/abt_lower.cpp b/src/mongo/db/exec/sbe/abt/abt_lower.cpp index f32fe8da661..a706d150d61 100644 --- a/src/mongo/db/exec/sbe/abt/abt_lower.cpp +++ b/src/mongo/db/exec/sbe/abt/abt_lower.cpp @@ -1008,45 +1008,22 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const IndexScanNode& n, co const PlanNodeId planNodeId = _nodeToGroupPropsMap.at(&n)._planNodeId; auto projectForKeyStringBounds = sbe::makeS<sbe::LimitSkipStage>( sbe::makeS<sbe::CoScanStage>(planNodeId), 1, boost::none, planNodeId); - if (hasLowerBound) { - seekKeySlotLower = _slotIdGenerator.generate(); - correlatedSlotsForJoin.push_back(seekKeySlotLower.value()); - projectForKeyStringBounds = sbe::makeProjectStage(std::move(projectForKeyStringBounds), - planNodeId, - seekKeySlotLower.value(), - std::move(lowerBoundExpr)); - } - if (hasUpperBound) { - seekKeySlotUpper = _slotIdGenerator.generate(); - correlatedSlotsForJoin.push_back(seekKeySlotUpper.value()); - projectForKeyStringBounds = sbe::makeProjectStage(std::move(projectForKeyStringBounds), - planNodeId, - seekKeySlotUpper.value(), - std::move(upperBoundExpr)); - } // Unused. boost::optional<sbe::value::SlotId> resultSlot; - auto result = sbe::makeS<sbe::IndexScanStage>(nss.uuid().get(), - indexDefName, - !indexSpec.isReverseOrder(), - resultSlot, - ridSlot, - boost::none, - indexKeysToInclude, - vars, - seekKeySlotLower, - seekKeySlotUpper, - nullptr /*yieldPolicy*/, - planNodeId); - - return sbe::makeS<sbe::LoopJoinStage>(std::move(projectForKeyStringBounds), - std::move(result), - sbe::makeSV(), - std::move(correlatedSlotsForJoin), - nullptr, - planNodeId); + return sbe::makeS<sbe::IndexScanStage>(nss.uuid().get(), + indexDefName, + !indexSpec.isReverseOrder(), + resultSlot, + ridSlot, + boost::none, + indexKeysToInclude, + vars, + std::move(lowerBoundExpr), + std::move(upperBoundExpr), + nullptr /*yieldPolicy*/, + planNodeId); } std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const SeekNode& n, diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp index ba8cd9704e4..520e68fe074 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp +++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp @@ -47,11 +47,11 @@ IndexScanStage::IndexScanStage(UUID collUuid, boost::optional<value::SlotId> snapshotIdSlot, IndexKeysInclusionSet indexKeysToInclude, value::SlotVector vars, - boost::optional<value::SlotId> seekKeySlotLow, - boost::optional<value::SlotId> seekKeySlotHigh, + std::unique_ptr<EExpression> seekKeyLow, + std::unique_ptr<EExpression> seekKeyHigh, PlanYieldPolicy* yieldPolicy, PlanNodeId nodeId) - : PlanStage(seekKeySlotLow ? "ixseek"_sd : "ixscan"_sd, yieldPolicy, nodeId), + : PlanStage(seekKeyLow ? "ixseek"_sd : "ixscan"_sd, yieldPolicy, nodeId), _collUuid(collUuid), _indexName(indexName), _forward(forward), @@ -60,40 +60,11 @@ IndexScanStage::IndexScanStage(UUID collUuid, _snapshotIdSlot(snapshotIdSlot), _indexKeysToInclude(indexKeysToInclude), _vars(std::move(vars)), - _seekKeySlotLow(seekKeySlotLow), - _seekKeySlotHigh(seekKeySlotHigh) { + _seekKeyLow(std::move(seekKeyLow)), + _seekKeyHigh(std::move(seekKeyHigh)) { // The valid state is when both boundaries, or none is set, or only low key is set. - invariant((_seekKeySlotLow && _seekKeySlotHigh) || (!_seekKeySlotLow && !_seekKeySlotHigh) || - (_seekKeySlotLow && !_seekKeySlotHigh)); - - invariant(_indexKeysToInclude.count() == _vars.size()); -} -IndexScanStage::IndexScanStage(UUID collUuid, - StringData indexName, - bool forward, - boost::optional<value::SlotId> recordSlot, - boost::optional<value::SlotId> recordIdSlot, - boost::optional<value::SlotId> snapshotIdSlot, - IndexKeysInclusionSet indexKeysToInclude, - value::SlotVector vars, - std::unique_ptr<EExpression> seekKeyLowVar, - std::unique_ptr<EExpression> seekKeyHighVar, - PlanYieldPolicy* yieldPolicy, - PlanNodeId nodeId) - : PlanStage(seekKeyLowVar ? "ixseek"_sd : "ixscan"_sd, yieldPolicy, nodeId), - _collUuid(collUuid), - _indexName(indexName), - _forward(forward), - _recordSlot(recordSlot), - _recordIdSlot(recordIdSlot), - _snapshotIdSlot(snapshotIdSlot), - _indexKeysToInclude(indexKeysToInclude), - _vars(std::move(vars)), - _seekKeyLowVar(std::move(seekKeyLowVar)), - _seekKeyHighVar(std::move(seekKeyHighVar)) { - // The valid state is when both boundaries, or none is set, or only low key is set. - invariant((_seekKeyLowVar && _seekKeyHighVar) || (!_seekKeyLowVar && !_seekKeyHighVar) || - (_seekKeyLowVar && !_seekKeyHighVar)); + invariant((_seekKeyLow && _seekKeyHigh) || (!_seekKeyLow && !_seekKeyHigh) || + (_seekKeyLow && !_seekKeyHigh)); invariant(_indexKeysToInclude.count() == _vars.size()); } @@ -107,8 +78,8 @@ std::unique_ptr<PlanStage> IndexScanStage::clone() const { _snapshotIdSlot, _indexKeysToInclude, _vars, - _seekKeySlotLow, - _seekKeySlotHigh, + _seekKeyLow->clone(), + _seekKeyHigh->clone(), _yieldPolicy, _commonStats.nodeId); } @@ -132,11 +103,13 @@ void IndexScanStage::prepare(CompileCtx& ctx) { uassert(4822821, str::stream() << "duplicate slot: " << _vars[idx], inserted); } - if (_seekKeySlotLow) { - _seekKeyLowAccessor = ctx.getAccessor(*_seekKeySlotLow); + if (_seekKeyLow) { + ctx.root = this; + _seekKeyLowCodes = _seekKeyLow->compile(ctx); } - if (_seekKeySlotHigh) { - _seekKeyHiAccessor = ctx.getAccessor(*_seekKeySlotHigh); + if (_seekKeyHigh) { + ctx.root = this; + _seekKeyHighCodes = _seekKeyHigh->compile(ctx); _seekKeyHighHolder = std::make_unique<value::OwnedValueAccessor>(); } _seekKeyLowHolder = std::make_unique<value::OwnedValueAccessor>(); @@ -162,16 +135,6 @@ void IndexScanStage::prepare(CompileCtx& ctx) { value::TypeTags::NumberInt64, value::bitcastFrom<uint64_t>(_opCtx->recoveryUnit()->getSnapshotId().toNumber())); } - - if (_seekKeyLowVar) { - ctx.root = this; - _seekKeyLowCodes = _seekKeyLowVar->compile(ctx); - } - if (_seekKeyHighVar) { - ctx.root = this; - _seekKeyHighCodes = _seekKeyHighVar->compile(ctx); - _seekKeyHighHolder = std::make_unique<value::OwnedValueAccessor>(); - } } value::SlotAccessor* IndexScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot) { @@ -317,34 +280,28 @@ void IndexScanStage::open(bool reOpen) { _cursor = entry->accessMethod()->asSortedData()->newCursor(_opCtx, _forward); } - if (_seekKeyLowAccessor && _seekKeyHiAccessor) { - auto [tagLow, valLow] = _seekKeyLowAccessor->getViewOfValue(); + if (_seekKeyLow && _seekKeyHigh) { + auto [ownedLow, tagLow, valLow] = _bytecode.run(_seekKeyLowCodes.get()); const auto msgTagLow = tagLow; uassert(4822851, str::stream() << "seek key is wrong type: " << msgTagLow, tagLow == value::TypeTags::ksValue); - _seekKeyLowHolder->reset(false, tagLow, valLow); + _seekKeyLowHolder->reset(ownedLow, tagLow, valLow); - auto [tagHi, valHi] = _seekKeyHiAccessor->getViewOfValue(); + auto [ownedHi, tagHi, valHi] = _bytecode.run(_seekKeyHighCodes.get()); const auto msgTagHi = tagHi; uassert(4822852, str::stream() << "seek key is wrong type: " << msgTagHi, tagHi == value::TypeTags::ksValue); - _seekKeyHighHolder->reset(false, tagHi, valHi); - } else if (_seekKeyLowAccessor) { - auto [tagLow, valLow] = _seekKeyLowAccessor->getViewOfValue(); + _seekKeyHighHolder->reset(ownedHi, tagHi, valHi); + } else if (_seekKeyLow) { + auto [ownedLow, tagLow, valLow] = _bytecode.run(_seekKeyLowCodes.get()); const auto msgTagLow = tagLow; uassert(4822853, str::stream() << "seek key is wrong type: " << msgTagLow, tagLow == value::TypeTags::ksValue); - _seekKeyLowHolder->reset(false, tagLow, valLow); - } else if (_seekKeyLowVar && _seekKeyHighVar) { - auto [ownedLow, tagLow, valLow] = _bytecode.run(_seekKeyLowCodes.get()); _seekKeyLowHolder->reset(ownedLow, tagLow, valLow); - - auto [ownedHigh, tagHigh, valHigh] = _bytecode.run(_seekKeyHighCodes.get()); - _seekKeyHighHolder->reset(ownedHigh, tagHigh, valHigh); } else { auto sdi = entry->accessMethod()->asSortedData()->getSortedDataInterface(); KeyString::Builder kb(sdi->getKeyStringVersion(), @@ -454,6 +411,7 @@ std::unique_ptr<PlanStageStats> IndexScanStage::getStats(bool includeDebugInfo) ret->specific = std::make_unique<IndexScanStats>(_specificStats); if (includeDebugInfo) { + DebugPrinter printer; BSONObjBuilder bob; bob.append("indexName", _indexName); bob.appendNumber("keysExamined", static_cast<long long>(_specificStats.keysExamined)); @@ -468,11 +426,11 @@ std::unique_ptr<PlanStageStats> IndexScanStage::getStats(bool includeDebugInfo) if (_snapshotIdSlot) { bob.appendNumber("snapshotIdSlot", static_cast<long long>(*_snapshotIdSlot)); } - if (_seekKeySlotLow) { - bob.appendNumber("seekKeySlotLow", static_cast<long long>(*_seekKeySlotLow)); + if (_seekKeyLow) { + bob.append("seekKeyLow", printer.print(_seekKeyLow->debugPrint())); } - if (_seekKeySlotHigh) { - bob.appendNumber("seekKeySlotHigh", static_cast<long long>(*_seekKeySlotHigh)); + if (_seekKeyHigh) { + bob.append("seekKeyHigh", printer.print(_seekKeyHigh->debugPrint())); } bob.append("outputSlots", _vars.begin(), _vars.end()); bob.append("indexKeysToInclude", _indexKeysToInclude.to_string()); @@ -489,17 +447,10 @@ const SpecificStats* IndexScanStage::getSpecificStats() const { std::vector<DebugPrinter::Block> IndexScanStage::debugPrint() const { auto ret = PlanStage::debugPrint(); - if (_seekKeySlotLow) { - DebugPrinter::addIdentifier(ret, _seekKeySlotLow.get()); - if (_seekKeySlotHigh) { - DebugPrinter::addIdentifier(ret, _seekKeySlotHigh.get()); - } else { - DebugPrinter::addIdentifier(ret, DebugPrinter::kNoneKeyword); - } - } else if (_seekKeyLowVar) { - DebugPrinter::addBlocks(ret, _seekKeyLowVar->debugPrint()); - if (_seekKeyHighVar) { - DebugPrinter::addBlocks(ret, _seekKeyHighVar->debugPrint()); + if (_seekKeyLow) { + DebugPrinter::addBlocks(ret, _seekKeyLow->debugPrint()); + if (_seekKeyHigh) { + DebugPrinter::addBlocks(ret, _seekKeyHigh->debugPrint()); } else { DebugPrinter::addIdentifier(ret, DebugPrinter::kNoneKeyword); } diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h index 3e61a086670..c57389b3434 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.h +++ b/src/mongo/db/exec/sbe/stages/ix_scan.h @@ -41,14 +41,14 @@ namespace mongo::sbe { /** * A stage that iterates the entries of a collection index, starting from a bound specified by the - * value in 'seekKeySlotLow' and ending (via IS_EOF) with the 'seekKeySlotHigh' bound. (An - * unspecified 'seekKeySlotHigh' scans to the end of the index. Leaving both bounds unspecified + * value in 'seekKeyLow' and ending (via IS_EOF) with the 'seekKeyHigh' bound. (An + * unspecified 'seekKeyHigh' scans to the end of the index. Leaving both bounds unspecified * scans the index from beginning to end.) * - * The input 'seekKeySlotLow' and 'seekKeySlotHigh' slots get read as part of the open (or re-open) - * call. A common use case for an IndexScanStage is to place it as the inner child of LoopJoinStage. - * The outer side of the LoopJoinStage determines the bounds, and the inner IndexScanStage iterates - * through all the entries within those bounds. + * The input 'seekKeyLow' and 'seekKeyHigh' EExpressions get evaluated as part of the open + * (or re-open) call. A common use case for an IndexScanStage is to place it as the inner child of + * LoopJoinStage. The outer side of the LoopJoinStage determines the bounds, and the inner + * IndexScanStage iterates through all the entries within those bounds. * * The "output" slots are * - 'recordSlot': the "KeyString" representing the index entry, @@ -82,21 +82,8 @@ public: boost::optional<value::SlotId> snapshotIdSlot, IndexKeysInclusionSet indexKeysToInclude, value::SlotVector vars, - boost::optional<value::SlotId> seekKeySlotLow, - boost::optional<value::SlotId> seekKeySlotHigh, - PlanYieldPolicy* yieldPolicy, - PlanNodeId nodeId); - - IndexScanStage(UUID collUuid, - StringData indexName, - bool forward, - boost::optional<value::SlotId> recordSlot, - boost::optional<value::SlotId> recordIdSlot, - boost::optional<value::SlotId> snapshotIdSlot, - IndexKeysInclusionSet indexKeysToInclude, - value::SlotVector vars, - std::unique_ptr<EExpression> seekKeyLowVar, - std::unique_ptr<EExpression> seekKeyHighVar, + std::unique_ptr<EExpression> seekKeyLow, + std::unique_ptr<EExpression> seekKeyHigh, PlanYieldPolicy* yieldPolicy, PlanNodeId nodeId); @@ -142,19 +129,16 @@ private: const boost::optional<value::SlotId> _snapshotIdSlot; const IndexKeysInclusionSet _indexKeysToInclude; const value::SlotVector _vars; - const boost::optional<value::SlotId> _seekKeySlotLow; - const boost::optional<value::SlotId> _seekKeySlotHigh; - std::unique_ptr<EExpression> _seekKeyLowVar; - std::unique_ptr<EExpression> _seekKeyHighVar; + std::unique_ptr<EExpression> _seekKeyLow; + std::unique_ptr<EExpression> _seekKeyHigh; - // For the EExpression overload for seek keys. + // Carries the compiled bytecode for the above '_seekKeyLow' and '_seekKeyHigh'. std::unique_ptr<vm::CodeFragment> _seekKeyLowCodes; std::unique_ptr<vm::CodeFragment> _seekKeyHighCodes; vm::ByteCode _bytecode; - // These members are default constructed to boost::none and are initialized when 'prepare()' // is called. Once they are set, they are never modified again. boost::optional<NamespaceString> _collName; diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp index 13ff93f7b53..ec8376d8095 100644 --- a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp +++ b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp @@ -318,8 +318,8 @@ generateOptimizedMultiIntervalIndexScan(StageBuilderState& state, indexSnapshotSlot, indexKeysToInclude, std::move(indexKeySlots), - lowKeySlot, - highKeySlot, + makeVariable(lowKeySlot), + makeVariable(highKeySlot), yieldPolicy, planNodeId); @@ -453,8 +453,8 @@ makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection, snapshotIdSlot, indexKeysToInclude, std::move(savedIndexKeySlots), - lowKeySlot, - boost::none, + makeVariable(lowKeySlot), + nullptr /* seekKeyHigh */, yieldPolicy, planNodeId); @@ -873,22 +873,30 @@ generateSingleIntervalIndexScan(StageBuilderState& state, // are present. Otherwise the low and high keys will be obtained via variable references to // runtime environment slots. sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projects; - auto makeKeySlot = [&](std::unique_ptr<KeyString::Value> key) { + auto makeKeyExpr = [&](std::unique_ptr<KeyString::Value> key) + -> std::pair<std::unique_ptr<sbe::EExpression>, boost::optional<sbe::value::SlotId>> { if (key) { - auto keySlot = slotIdGenerator->generate(); - projects.emplace( - keySlot, + return std::pair{ makeConstant(sbe::value::TypeTags::ksValue, - sbe::value::bitcastFrom<KeyString::Value*>(key.release()))); - return keySlot; + sbe::value::bitcastFrom<KeyString::Value*>(key.release())), + boost::none}; } else { auto keySlot = state.data->env->registerSlot( sbe::value::TypeTags::Nothing, 0, true /* owned */, slotIdGenerator); - return keySlot; + return std::pair{makeVariable(keySlot), keySlot}; } }; - auto lowKeySlot = makeKeySlot(std::move(lowKey)); - auto highKeySlot = makeKeySlot(std::move(highKey)); + + std::unique_ptr<sbe::EExpression> lowKeyExpr; + std::unique_ptr<sbe::EExpression> highKeyExpr; + boost::optional<sbe::value::SlotId> lowKeySlot; + boost::optional<sbe::value::SlotId> highKeySlot; + std::tie(lowKeyExpr, lowKeySlot) = makeKeyExpr(std::move(lowKey)); + std::tie(highKeyExpr, highKeySlot) = makeKeyExpr(std::move(highKey)); + tassert(0000000, + "Either both lowKeySlot and highKeySlot will exist or none of them will exist", + (lowKeySlot && highKeySlot) || (!lowKeySlot && !highKeySlot)); + if (indexIdSlot) { // Construct a copy of 'indexName' to project for use in the index consistency check. @@ -924,8 +932,8 @@ generateSingleIntervalIndexScan(StageBuilderState& state, return sbe::makeS<sbe::FilterStage</* IsConst */ true, /* IsEof */ false>>( std::move(childStage), makeBinaryOp(sbe::EPrimBinary::logicAnd, - makeFunction("exists", makeVariable(lowKeySlot)), - makeFunction("exists", makeVariable(highKeySlot))), + makeFunction("exists", lowKeyExpr->clone()), + makeFunction("exists", highKeyExpr->clone())), planNodeId); }(); @@ -948,8 +956,8 @@ generateSingleIntervalIndexScan(StageBuilderState& state, indexSnapshotSlot, indexKeysToInclude, std::move(indexKeySlots), - lowKeySlot, - highKeySlot, + std::move(lowKeyExpr), + std::move(highKeyExpr), yieldPolicy, planNodeId); @@ -975,11 +983,11 @@ generateSingleIntervalIndexScan(StageBuilderState& state, sbe::makeS<sbe::LoopJoinStage>(std::move(lowHighKeyBranch), std::move(stage), std::move(outerSv), - sbe::makeSV(lowKeySlot, highKeySlot), + sbe::makeSV(), nullptr, planNodeId), boost::make_optional(shouldRegisterLowHighKeyInRuntimeEnv, - std::pair(lowKeySlot, highKeySlot))}; + std::pair(*lowKeySlot, *highKeySlot))}; } std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( diff --git a/src/mongo/db/query/sbe_stage_builder_lookup.cpp b/src/mongo/db/query/sbe_stage_builder_lookup.cpp index b1077b0ae79..4525bedd4f7 100644 --- a/src/mongo/db/query/sbe_stage_builder_lookup.cpp +++ b/src/mongo/db/query/sbe_stage_builder_lookup.cpp @@ -804,8 +804,8 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage( snapshotIdSlot, IndexKeysInclusionSet{} /* indexKeysToInclude */, makeSV() /* vars */, - lowKeySlot, - highKeySlot, + makeVariable(lowKeySlot), + makeVariable(highKeySlot), yieldPolicy, nodeId); |