diff options
author | Nikita Lapkov <nikita.lapkov@mongodb.com> | 2022-04-12 15:24:11 +0000 |
---|---|---|
committer | Evergreen Agent <no-reply@evergreen.mongodb.com> | 2022-04-12 22:01:38 +0000 |
commit | 3780d57e1e378b55a887a43115c8d303e8962a6b (patch) | |
tree | 5352040139db98e614b96434285676c59d72efaa /src/mongo/db/query/plan_executor_sbe.cpp | |
parent | a20188ca06dbd604f4386db97d4199124f551e74 (diff) | |
download | mongo-3780d57e1e378b55a887a43115c8d303e8962a6b.tar.gz |
SERVER-64597 Produce SBE Object instead of BSON in $lookup and $group
Diffstat (limited to 'src/mongo/db/query/plan_executor_sbe.cpp')
-rw-r--r-- | src/mongo/db/query/plan_executor_sbe.cpp | 246 |
1 files changed, 226 insertions, 20 deletions
diff --git a/src/mongo/db/query/plan_executor_sbe.cpp b/src/mongo/db/query/plan_executor_sbe.cpp index 0861ef6b35e..c8d516cb718 100644 --- a/src/mongo/db/query/plan_executor_sbe.cpp +++ b/src/mongo/db/query/plan_executor_sbe.cpp @@ -139,7 +139,7 @@ void PlanExecutorSBE::saveState() { } _yieldPolicy->setYieldable(nullptr); - _lastGetNext = {}; + _lastGetNext = BSONObj(); } void PlanExecutorSBE::restoreState(const RestoreContext& context) { @@ -197,10 +197,10 @@ PlanExecutor::ExecState PlanExecutorSBE::getNextDocument(Document* objOut, Recor checkFailPointPlanExecAlwaysFails(); - BSONObj obj; - auto result = getNext(&obj, dlOut); + Document obj; + auto result = getNextImpl(&obj, dlOut); if (result == PlanExecutor::ExecState::ADVANCED) { - *objOut = Document{std::move(obj)}; + *objOut = std::move(obj); } return result; } @@ -210,16 +210,46 @@ PlanExecutor::ExecState PlanExecutorSBE::getNext(BSONObj* out, RecordId* dlOut) checkFailPointPlanExecAlwaysFails(); + BSONObj obj; + auto result = getNextImpl(&obj, dlOut); + if (result == PlanExecutor::ExecState::ADVANCED) { + *out = std::move(obj); + } + return result; +} + +template <typename ObjectType> +sbe::PlanState fetchNextImpl(sbe::PlanStage* root, + sbe::value::SlotAccessor* resultSlot, + sbe::value::SlotAccessor* recordIdSlot, + ObjectType* out, + RecordId* dlOut, + bool returnOwnedBson); + +template <typename ObjectType> +PlanExecutor::ExecState PlanExecutorSBE::getNextImpl(ObjectType* out, RecordId* dlOut) { + constexpr bool isDocument = std::is_same_v<ObjectType, Document>; + constexpr bool isBson = std::is_same_v<ObjectType, BSONObj>; + static_assert(isDocument || isBson); + + invariant(!_isDisposed); + + checkFailPointPlanExecAlwaysFails(); + if (!_stash.empty()) { auto&& [doc, recordId] = _stash.front(); - *out = std::move(doc); + if constexpr (isBson) { + *out = std::move(doc); + } else { + *out = Document{std::move(doc)}; + } if (dlOut && recordId) { *dlOut = *recordId; } _stash.pop_front(); return PlanExecutor::ExecState::ADVANCED; } else if (_root->getCommonStats()->isEOF) { - // If we had stashed elements and consumed them all, but the PlanStage has also + // If we had stashed elements and consumed them all, but the PlanStage is also // already exhausted, we can return EOF straight away. Otherwise, proceed with // fetching the next document. _root->close(); @@ -263,11 +293,11 @@ PlanExecutor::ExecState PlanExecutorSBE::getNext(BSONObj* out, RecordId* dlOut) invariant(_state == State::kOpened); auto result = - fetchNext(_root.get(), _result, _resultRecordId, out, dlOut, _mustReturnOwnedBson); + fetchNextImpl(_root.get(), _result, _resultRecordId, out, dlOut, _mustReturnOwnedBson); if (result == sbe::PlanState::IS_EOF) { _root->close(); _state = State::kClosed; - _lastGetNext = {}; + _lastGetNext = BSONObj(); if (MONGO_unlikely(planExecutorHangBeforeShouldWaitForInserts.shouldFail( [this](const BSONObj& data) { @@ -298,12 +328,32 @@ PlanExecutor::ExecState PlanExecutorSBE::getNext(BSONObj* out, RecordId* dlOut) invariant(result == sbe::PlanState::ADVANCED); if (_mustReturnOwnedBson) { - _lastGetNext = *out; + if constexpr (isBson) { + _lastGetNext = *out; + } else { + auto bson = out->toBsonIfTriviallyConvertible(); + if (bson) { + // This basically means that the 'Document' is just a wrapper around BSON + // returned by the plan. In this case, 'out' must own it. + invariant(out->isOwned()); + _lastGetNext = *bson; + } else { + // 'Document' was created from 'sbe::Object' and there is no backing BSON for + // it. + _lastGetNext = BSONObj(); + } + } } return PlanExecutor::ExecState::ADVANCED; } } +// Explicitly instantiate the only 2 types supported by 'PlanExecutorSBE::getNextImpl'. +template PlanExecutor::ExecState PlanExecutorSBE::getNextImpl<BSONObj>(BSONObj* out, + RecordId* dlOut); +template PlanExecutor::ExecState PlanExecutorSBE::getNextImpl<Document>(Document* out, + RecordId* dlOut); + Timestamp PlanExecutorSBE::getLatestOplogTimestamp() const { if (_rootData.shouldTrackLatestOplogTimestamp) { tassert(5567201, @@ -350,12 +400,132 @@ BSONObj PlanExecutorSBE::getPostBatchResumeToken() const { return {}; } -sbe::PlanState fetchNext(sbe::PlanStage* root, - sbe::value::SlotAccessor* resultSlot, - sbe::value::SlotAccessor* recordIdSlot, - BSONObj* out, - RecordId* dlOut, - bool returnOwnedBson) { +namespace { + +Document convertToDocument(const sbe::value::Object& obj); + +Value convertToValue(sbe::value::TypeTags tag, sbe::value::Value val) { + switch (tag) { + case sbe::value::TypeTags::Nothing: + return Value(); + + case sbe::value::TypeTags::NumberInt32: + return Value(sbe::value::bitcastTo<int32_t>(val)); + + case sbe::value::TypeTags::NumberInt64: + return Value(sbe::value::bitcastTo<long long>(val)); + + case sbe::value::TypeTags::NumberDouble: + return Value(sbe::value::bitcastTo<double>(val)); + + case sbe::value::TypeTags::NumberDecimal: + return Value(sbe::value::bitcastTo<Decimal128>(val)); + + case sbe::value::TypeTags::Date: + return Value(Date_t::fromMillisSinceEpoch(sbe::value::bitcastTo<int64_t>(val))); + + case sbe::value::TypeTags::Timestamp: + return Value(Timestamp(sbe::value::bitcastTo<uint64_t>(val))); + + case sbe::value::TypeTags::Boolean: + return Value(sbe::value::bitcastTo<bool>(val)); + + case sbe::value::TypeTags::Null: + return Value(BSONNULL); + + case sbe::value::TypeTags::StringSmall: + case sbe::value::TypeTags::StringBig: + case sbe::value::TypeTags::bsonString: + case sbe::value::TypeTags::bsonSymbol: + return Value(sbe::value::getStringOrSymbolView(tag, val)); + + case sbe::value::TypeTags::Array: + case sbe::value::TypeTags::ArraySet: { + std::vector<Value> vals; + auto enumerator = sbe::value::ArrayEnumerator{tag, val}; + while (!enumerator.atEnd()) { + auto [arrTag, arrVal] = enumerator.getViewOfValue(); + enumerator.advance(); + vals.push_back(convertToValue(arrTag, arrVal)); + } + return Value(std::move(vals)); + } + + case sbe::value::TypeTags::Object: + return Value(convertToDocument(*sbe::value::getObjectView(val))); + + case sbe::value::TypeTags::ObjectId: + return Value(OID::from(sbe::value::getObjectIdView(val)->data())); + + case sbe::value::TypeTags::MinKey: + return Value(kMinBSONKey); + + case sbe::value::TypeTags::MaxKey: + return Value(kMaxBSONKey); + + case sbe::value::TypeTags::bsonObject: + return Value(BSONObj{sbe::value::bitcastTo<const char*>(val)}); + + case sbe::value::TypeTags::bsonArray: + return Value(BSONArray{BSONObj{sbe::value::bitcastTo<const char*>(val)}}); + + case sbe::value::TypeTags::bsonObjectId: + return Value(OID::from(sbe::value::bitcastTo<const char*>(val))); + + case sbe::value::TypeTags::bsonBinData: + return Value(BSONBinData(sbe::value::getBSONBinData(tag, val), + sbe::value::getBSONBinDataSize(tag, val), + sbe::value::getBSONBinDataSubtype(tag, val))); + + case sbe::value::TypeTags::bsonUndefined: + return Value(BSONUndefined); + + case sbe::value::TypeTags::bsonRegex: { + auto regex = sbe::value::getBsonRegexView(val); + return Value(BSONRegEx(regex.pattern, regex.flags)); + } + + case sbe::value::TypeTags::bsonJavascript: + return Value(sbe::value::getBsonJavascriptView(val)); + + case sbe::value::TypeTags::bsonDBPointer: { + auto dbptr = sbe::value::getBsonDBPointerView(val); + return Value(BSONDBRef(dbptr.ns, OID::from(dbptr.id))); + } + + case sbe::value::TypeTags::bsonCodeWScope: { + auto bcws = sbe::value::getBsonCodeWScopeView(val); + return Value(BSONCodeWScope(bcws.code, BSONObj(bcws.scope))); + } + + default: + MONGO_UNREACHABLE; + } +} + +Document convertToDocument(const sbe::value::Object& obj) { + MutableDocument doc; + for (size_t idx = 0; idx < obj.size(); ++idx) { + auto [tag, val] = obj.getAt(idx); + const auto& name = obj.field(idx); + doc.addField(name, convertToValue(tag, val)); + } + return doc.freeze(); +} + +} // namespace + +template <typename ObjectType> +sbe::PlanState fetchNextImpl(sbe::PlanStage* root, + sbe::value::SlotAccessor* resultSlot, + sbe::value::SlotAccessor* recordIdSlot, + ObjectType* out, + RecordId* dlOut, + bool returnOwnedBson) { + constexpr bool isDocument = std::is_same_v<ObjectType, Document>; + constexpr bool isBson = std::is_same_v<ObjectType, BSONObj>; + static_assert(isDocument || isBson); + invariant(out); auto state = root->getNext(); @@ -371,17 +541,28 @@ sbe::PlanState fetchNext(sbe::PlanStage* root, if (resultSlot) { auto [tag, val] = resultSlot->getViewOfValue(); if (tag == sbe::value::TypeTags::Object) { - BSONObjBuilder bb; - sbe::bson::convertToBsonObj(bb, sbe::value::getObjectView(val)); - *out = bb.obj(); + if constexpr (isBson) { + BSONObjBuilder bb; + sbe::bson::convertToBsonObj(bb, sbe::value::getObjectView(val)); + *out = bb.obj(); + } else { + *out = convertToDocument(*sbe::value::getObjectView(val)); + } } else if (tag == sbe::value::TypeTags::bsonObject) { + BSONObj result; if (returnOwnedBson) { auto [ownedTag, ownedVal] = resultSlot->copyOrMoveValue(); auto sharedBuf = SharedBuffer(UniqueBuffer::reclaim(sbe::value::bitcastTo<char*>(ownedVal))); - *out = BSONObj(std::move(sharedBuf)); + result = BSONObj{std::move(sharedBuf)}; } else { - *out = BSONObj(sbe::value::bitcastTo<const char*>(val)); + result = BSONObj{sbe::value::bitcastTo<const char*>(val)}; + } + + if constexpr (isBson) { + *out = std::move(result); + } else { + *out = Document{std::move(result)}; } } else { // The query is supposed to return an object. @@ -398,4 +579,29 @@ sbe::PlanState fetchNext(sbe::PlanStage* root, } return state; } + +template sbe::PlanState fetchNextImpl<BSONObj>(sbe::PlanStage* root, + sbe::value::SlotAccessor* resultSlot, + sbe::value::SlotAccessor* recordIdSlot, + BSONObj* out, + RecordId* dlOut, + bool returnOwnedBson); + +template sbe::PlanState fetchNextImpl<Document>(sbe::PlanStage* root, + sbe::value::SlotAccessor* resultSlot, + sbe::value::SlotAccessor* recordIdSlot, + Document* out, + RecordId* dlOut, + bool returnOwnedBson); + +// NOTE: We intentionally do not expose overload for the 'Document' type. The only interface to get +// result from plan in 'Document' type is to call 'PlanExecutorSBE::getNextDocument()'. +sbe::PlanState fetchNext(sbe::PlanStage* root, + sbe::value::SlotAccessor* resultSlot, + sbe::value::SlotAccessor* recordIdSlot, + BSONObj* out, + RecordId* dlOut, + bool returnOwnedBson) { + return fetchNextImpl(root, resultSlot, recordIdSlot, out, dlOut, returnOwnedBson); +} } // namespace mongo |