diff options
author | Alexander Ignatyev <alexander.ignatyev@mongodb.com> | 2021-09-17 08:51:15 +0000 |
---|---|---|
committer | Evergreen Agent <no-reply@evergreen.mongodb.com> | 2021-09-17 09:29:18 +0000 |
commit | 2a82df1a0433d21bfb55ca7e6ff2b0d47822d34c (patch) | |
tree | 0de92bc8be9e07b9945afc511f873013014372b5 /src | |
parent | f508575e0a4f72979247176c279c5497da0a260a (diff) | |
download | mongo-2a82df1a0433d21bfb55ca7e6ff2b0d47822d34c.tar.gz |
SERVER-59331 Add PlanCache's and EExpression's getCompileTimeSize() functions to calculate compile-time size of SBE Plan
Diffstat (limited to 'src')
52 files changed, 916 insertions, 2 deletions
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript index f7ab8543480..68c00419ae6 100644 --- a/src/mongo/db/exec/sbe/SConscript +++ b/src/mongo/db/exec/sbe/SConscript @@ -35,6 +35,7 @@ sbeEnv.Library( target='query_sbe', source=[ 'expressions/expression.cpp', + 'size_estimator.cpp', 'stages/branch.cpp', 'stages/bson_scan.cpp', 'stages/check_bounds.cpp', @@ -155,6 +156,7 @@ env.CppUnitTest( 'sbe_merge_join_test.cpp', 'sbe_mkobj_test.cpp', 'sbe_numeric_convert_test.cpp', + 'sbe_plan_size_test.cpp', 'sbe_sort_test.cpp', 'sbe_sorted_merge_test.cpp', 'sbe_spool_test.cpp', diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp index d939053b9fb..ef629ba7c02 100644 --- a/src/mongo/db/exec/sbe/expressions/expression.cpp +++ b/src/mongo/db/exec/sbe/expressions/expression.cpp @@ -34,6 +34,7 @@ #include <iomanip> #include <sstream> +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/sbe/stages/spool.h" #include "mongo/db/exec/sbe/stages/stages.h" #include "mongo/util/str.h" @@ -83,6 +84,13 @@ std::vector<DebugPrinter::Block> EConstant::debugPrint() const { return ret; } +size_t EConstant::estimateSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_tag, _val); + size += size_estimator::estimate(_nodes); + return size; +} + std::unique_ptr<EExpression> EVariable::clone() const { return _frameId ? std::make_unique<EVariable>(*_frameId, _var) : std::make_unique<EVariable>(_var); @@ -282,6 +290,11 @@ std::vector<DebugPrinter::Block> EPrimBinary::debugPrint() const { return ret; } +size_t EPrimBinary::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_nodes); +} + + std::unique_ptr<EExpression> EPrimUnary::clone() const { return std::make_unique<EPrimUnary>(_op, _nodes[0]->clone()); } @@ -321,6 +334,10 @@ std::vector<DebugPrinter::Block> EPrimUnary::debugPrint() const { return ret; } +size_t EPrimUnary::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_nodes); +} + std::unique_ptr<EExpression> EFunction::clone() const { Vector args; args.reserve(_nodes.size()); @@ -611,6 +628,10 @@ std::vector<DebugPrinter::Block> EFunction::debugPrint() const { return ret; } +size_t EFunction::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_name) + size_estimator::estimate(_nodes); +} + std::unique_ptr<EExpression> EIf::clone() const { return std::make_unique<EIf>(_nodes[0]->clone(), _nodes[1]->clone(), _nodes[2]->clone()); } @@ -659,6 +680,10 @@ std::vector<DebugPrinter::Block> EIf::debugPrint() const { return ret; } +size_t EIf::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_nodes); +} + std::unique_ptr<EExpression> ELocalBind::clone() const { Vector binds; binds.reserve(_nodes.size() - 1); @@ -712,6 +737,10 @@ std::vector<DebugPrinter::Block> ELocalBind::debugPrint() const { return ret; } +size_t ELocalBind::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_nodes); +} + std::unique_ptr<EExpression> ELocalLambda::clone() const { return std::make_unique<ELocalLambda>(_frameId, _nodes.back()->clone()); } @@ -751,6 +780,10 @@ std::vector<DebugPrinter::Block> ELocalLambda::debugPrint() const { return ret; } +size_t ELocalLambda::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_nodes); +} + std::unique_ptr<EExpression> EFail::clone() const { return std::make_unique<EFail>(_code, getStringView(_messageTag, _messageVal)); @@ -784,6 +817,11 @@ std::vector<DebugPrinter::Block> EFail::debugPrint() const { return ret; } +size_t EFail::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_messageTag, _messageVal) + + size_estimator::estimate(_nodes); +} + std::unique_ptr<EExpression> ENumericConvert::clone() const { return std::make_unique<ENumericConvert>(_nodes[0]->clone(), _target); } @@ -827,6 +865,10 @@ std::vector<DebugPrinter::Block> ENumericConvert::debugPrint() const { return ret; } +size_t ENumericConvert::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_nodes); +} + std::unique_ptr<EExpression> ETypeMatch::clone() const { return std::make_unique<ETypeMatch>(_nodes[0]->clone(), _typeMask); } @@ -856,6 +898,11 @@ std::vector<DebugPrinter::Block> ETypeMatch::debugPrint() const { return ret; } +size_t ETypeMatch::estimateSize() const { + return sizeof(*this) + size_estimator::estimate(_nodes); +} + + RuntimeEnvironment::RuntimeEnvironment(const RuntimeEnvironment& other) : _state{other._state}, _isSmp{other._isSmp} { for (auto&& [slotId, index] : _state->slots) { diff --git a/src/mongo/db/exec/sbe/expressions/expression.h b/src/mongo/db/exec/sbe/expressions/expression.h index 8ecb3ecf737..ae418eae99b 100644 --- a/src/mongo/db/exec/sbe/expressions/expression.h +++ b/src/mongo/db/exec/sbe/expressions/expression.h @@ -286,6 +286,11 @@ public: virtual std::vector<DebugPrinter::Block> debugPrint() const = 0; + /** + * Estimates the size of the current expression node and its children. + */ + virtual size_t estimateSize() const = 0; + protected: Vector _nodes; @@ -370,6 +375,8 @@ public: vm::CodeFragment compileDirect(CompileCtx& ctx) const override; std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: value::TypeTags _tag; @@ -392,6 +399,9 @@ public: vm::CodeFragment compileDirect(CompileCtx& ctx) const override; std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final { + return sizeof(*this); + } private: value::SlotId _var; @@ -457,6 +467,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: Op _op; }; @@ -482,6 +494,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: Op _op; }; @@ -504,6 +518,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: std::string _name; }; @@ -527,6 +543,8 @@ public: vm::CodeFragment compileDirect(CompileCtx& ctx) const override; std::vector<DebugPrinter::Block> debugPrint() const override; + + size_t estimateSize() const final; }; /** @@ -547,6 +565,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: FrameId _frameId; }; @@ -567,6 +587,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: FrameId _frameId; }; @@ -590,6 +612,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: ErrorCodes::Error _code; value::TypeTags _messageTag; @@ -622,6 +646,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: value::TypeTags _target; }; @@ -644,6 +670,8 @@ public: std::vector<DebugPrinter::Block> debugPrint() const override; + size_t estimateSize() const final; + private: uint32_t _typeMask; }; diff --git a/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp b/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp new file mode 100644 index 00000000000..aa6ad8545ee --- /dev/null +++ b/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp @@ -0,0 +1,321 @@ +/** + * Copyright (C) 2021-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. + */ + +#include "mongo/db/exec/sbe/stages/branch.h" +#include "mongo/db/exec/sbe/stages/bson_scan.h" +#include "mongo/db/exec/sbe/stages/check_bounds.h" +#include "mongo/db/exec/sbe/stages/co_scan.h" +#include "mongo/db/exec/sbe/stages/exchange.h" +#include "mongo/db/exec/sbe/stages/filter.h" +#include "mongo/db/exec/sbe/stages/hash_agg.h" +#include "mongo/db/exec/sbe/stages/hash_join.h" +#include "mongo/db/exec/sbe/stages/ix_scan.h" +#include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/db/exec/sbe/stages/makeobj.h" +#include "mongo/db/exec/sbe/stages/merge_join.h" +#include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/stages/scan.h" +#include "mongo/db/exec/sbe/stages/sort.h" +#include "mongo/db/exec/sbe/stages/sorted_merge.h" +#include "mongo/db/exec/sbe/stages/spool.h" +#include "mongo/db/exec/sbe/stages/traverse.h" +#include "mongo/db/exec/sbe/stages/union.h" +#include "mongo/db/exec/sbe/stages/unique.h" +#include "mongo/db/exec/sbe/stages/unwind.h" +#include "mongo/db/exec/sbe/values/slot.h" +#include "mongo/db/exec/sbe/values/value.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::sbe { + +class PlanSizeTest : public unittest::Test { +public: + void setUp() override { + _slotIdGenerator = std::make_unique<value::SlotIdGenerator>(); + } + + void tearDown() override { + _slotIdGenerator.reset(); + } + + value::SlotId generateSlotId() { + return _slotIdGenerator->generate(); + } + + static std::unique_ptr<PlanStage> mockS() { + return makeS<CoScanStage>(kEmptyPlanNodeId); + } + + value::SlotVector mockSV() { + return makeSV(generateSlotId()); + } + + std::unique_ptr<EExpression> mockE() { + return makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)); + } + + /** + * PlanSize is planform-dependent so here we just assert that the size is a reasonable number: + * bigger then zero and not too big. + * A too big number might mean unwanted wrapping around on unsigned integer. + */ + void assertPlanSize(const PlanStage& stage) { + size_t size = stage.estimateCompileTimeSize(); + ASSERT_LT(0ul, size); + ASSERT_GT(10000ul, size); + } + +private: + std::unique_ptr<value::SlotIdGenerator> _slotIdGenerator; +}; + +TEST_F(PlanSizeTest, Branch) { + auto stage = makeS<BranchStage>( + mockS(), mockS(), mockE(), mockSV(), mockSV(), mockSV(), kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, BsonScan) { + auto stage = makeS<BSONScanStage>(nullptr, + nullptr, + generateSlotId(), + std::vector<std::string>{2}, + mockSV(), + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, CheckBounds) { + CheckBoundsParams params{ + IndexBounds{}, BSONObj{}, int{}, KeyString::Version{}, Ordering::allAscending()}; + auto stage = makeS<CheckBoundsStage>( + mockS(), params, generateSlotId(), generateSlotId(), generateSlotId(), kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, CoScan) { + auto stage = makeS<CoScanStage>(kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Exchange) { + auto stage = makeS<ExchangeConsumer>( + mockS(), 1, makeSV(), ExchangePolicy::broadcast, nullptr, mockE(), kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Filter) { + auto stage = makeS<FilterStage<true>>(mockS(), mockE(), kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, HashAgg) { + auto stage = makeS<HashAggStage>(mockS(), + mockSV(), + makeEM(generateSlotId(), mockE()), + makeSV(), + true, + generateSlotId(), + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, HashJoin) { + auto stage = makeS<HashJoinStage>(mockS(), + mockS(), + mockSV(), + makeSV(), + mockSV(), + makeSV(), + generateSlotId(), + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, IndexScan) { + auto collUuid = CollectionUUID::parse("00000000-0000-0000-0000-000000000000").getValue(); + auto stage = makeS<IndexScanStage>(collUuid, + StringData(), + true, + generateSlotId(), + generateSlotId(), + generateSlotId(), + IndexKeysInclusionSet(1), + mockSV(), + generateSlotId(), + generateSlotId(), + nullptr, + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, LimitSkip) { + auto stage = makeS<LimitSkipStage>(mockS(), 200, 300, kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, LoopJoin) { + auto stage = + makeS<LoopJoinStage>(mockS(), mockS(), makeSV(), makeSV(), nullptr, kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, MakeObj) { + auto stage = makeS<MakeObjStage>(mockS(), + generateSlotId(), + generateSlotId(), + MakeObjFieldBehavior::keep, + std::vector<std::string>(), + std::vector<std::string>(), + makeSV(), + true, + false, + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, MergeJoin) { + std::vector<value::SortDirection> sortDirs(1, value::SortDirection::Ascending); + auto stage = makeS<MergeJoinStage>(mockS(), + mockS(), + mockSV(), + mockSV(), + mockSV(), + mockSV(), + std::move(sortDirs), + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Project) { + auto stage = makeProjectStage( + mockS(), kEmptyPlanNodeId, generateSlotId(), mockE(), generateSlotId(), mockE()); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Scan) { + auto collUuid = CollectionUUID::parse("00000000-0000-0000-0000-000000000000").getValue(); + auto stage = makeS<ScanStage>(collUuid, + generateSlotId(), + generateSlotId(), + generateSlotId(), + generateSlotId(), + generateSlotId(), + generateSlotId(), + boost::none, + std::vector<std::string>{"field"}, + mockSV(), + generateSlotId(), + true, + nullptr, + kEmptyPlanNodeId, + ScanCallbacks()); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Sort) { + auto stage = + makeS<SortStage>(mockS(), + mockSV(), + std::vector<value::SortDirection>{value::SortDirection::Ascending}, + mockSV(), + std::numeric_limits<std::size_t>::max(), + 204857600, + false, + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, SortedMerge) { + std::vector<value::SortDirection> sortDir{value::SortDirection::Ascending}; + const size_t numSlots = 4; + + PlanStage::Vector inputScans; + std::vector<value::SlotVector> inputSlots; + std::vector<value::SlotVector> inputKeys; + std::vector<value::SlotVector> inputVals; + for (size_t i = 0; i < numSlots; ++i) { + inputScans.push_back(mockS()); + inputKeys.push_back(mockSV()); + inputVals.push_back(mockSV()); + } + + auto stage = makeS<SortedMergeStage>(std::move(inputScans), + std::move(inputKeys), + std::move(sortDir), + std::move(inputVals), + mockSV(), + kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, SpoolLazyProducer) { + auto stage = makeS<SpoolLazyProducerStage>(mockS(), 1, mockSV(), nullptr, kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, SpoolConsumer) { + auto stage = makeS<SpoolConsumerStage<true>>(1, mockSV(), kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Traverse) { + auto stage = makeS<TraverseStage>(mockS(), + mockS(), + generateSlotId(), + generateSlotId(), + generateSlotId(), + mockSV(), + nullptr, + nullptr, + kEmptyPlanNodeId, + boost::none); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Union) { + auto scanStages = makeSs(mockS(), mockS()); + std::vector<value::SlotVector> scanInputVals{mockSV(), mockSV()}; + auto stage = makeS<UnionStage>( + std::move(scanStages), std::move(scanInputVals), mockSV(), kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Unique) { + auto stage = makeS<UniqueStage>(mockS(), mockSV(), kEmptyPlanNodeId); + assertPlanSize(*stage); +} + +TEST_F(PlanSizeTest, Unwind) { + auto stage = makeS<UnwindStage>( + mockS(), generateSlotId(), generateSlotId(), generateSlotId(), false, kEmptyPlanNodeId); + assertPlanSize(*stage); +} +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/size_estimator.cpp b/src/mongo/db/exec/sbe/size_estimator.cpp new file mode 100644 index 00000000000..0a11248287f --- /dev/null +++ b/src/mongo/db/exec/sbe/size_estimator.cpp @@ -0,0 +1,67 @@ +/** + * Copyright (C) 2021-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. + */ + +#include "mongo/db/exec/sbe/size_estimator.h" + +namespace mongo::sbe::size_estimator { + +size_t estimate(const IndexBounds& indexBounds) { + size_t size = estimate(indexBounds.startKey); + size += estimate(indexBounds.endKey); + size += estimate(indexBounds.fields); + return size; +} + +size_t estimate(const OrderedIntervalList& list) { + size_t size = estimate(list.name); + size += estimate(list.intervals); + return size; +} + +size_t estimate(const Interval& interval) { + size_t size = estimate(interval._intervalData); + size += estimate(interval.start); + size += estimate(interval.end); + return size; +} + +size_t estimate(const IndexSeekPoint& indexSeekPoint) { + size_t size = estimate(indexSeekPoint.keyPrefix); + size += estimate(indexSeekPoint.keySuffix); + size += estimate(indexSeekPoint.suffixInclusive); + return size; +} + +size_t estimate(const IndexBoundsChecker& checker) { + size_t size = sbe::size_estimator::estimate(checker._curInterval); + size += sbe::size_estimator::estimate(checker._expectedDirection); + size += sbe::size_estimator::estimate(checker._keyValues); + return size; +} +} // namespace mongo::sbe::size_estimator diff --git a/src/mongo/db/exec/sbe/size_estimator.h b/src/mongo/db/exec/sbe/size_estimator.h new file mode 100644 index 00000000000..b5dafdb68da --- /dev/null +++ b/src/mongo/db/exec/sbe/size_estimator.h @@ -0,0 +1,154 @@ +/** + * Copyright (C) 2021-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. + */ +#pragma once + +#include <absl/container/flat_hash_map.h> +#include <absl/container/inlined_vector.h> +#include <memory> +#include <string> +#include <vector> + +#include "mongo/bson/util/builder.h" +#include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/plan_stats.h" +#include "mongo/db/exec/sbe/stages/stages.h" +#include "mongo/db/exec/sbe/values/slot.h" +#include "mongo/db/query/index_bounds.h" +#include "mongo/db/storage/index_entry_comparison.h" + +/** + * Contains a set of functions for shallow estimating the size of allocated on the heap objects + * in a given container. The functions do not take account of the size of the given container, only + * its values allocated on the heap. + * They are used to calculate compile-time size of an SBE tree. + */ +namespace mongo::sbe::size_estimator { + +size_t estimate(const IndexBounds& indexBounds); +size_t estimate(const OrderedIntervalList& list); +size_t estimate(const Interval& interval); +size_t estimate(const IndexSeekPoint& indexSeekPoint); + +inline size_t estimate(const std::unique_ptr<PlanStage>& planStage) { + return planStage->estimateCompileTimeSize(); +} + +inline size_t estimate(const std::unique_ptr<EExpression>& expr) { + return expr->estimateSize(); +} + +inline size_t estimate(value::TypeTags tag, value::Value val) { + size_t size = value::getApproximateSize(tag, val); + return std::max(static_cast<size_t>(0), size - sizeof(tag) - sizeof(val)); +} + +inline size_t estimate(const std::string& str) { + return sizeof(std::string::value_type) * str.capacity(); +} + +inline size_t estimate(const BSONObj& bson) { + return bson.objsize(); +} + +size_t estimate(const IndexBoundsChecker& checker); + +// Calculate sizes of heap-allocated values only. Therefore, sizes of scalar values are always 0. +template <typename S, std::enable_if_t<std::is_scalar_v<S>, bool> = true> +inline size_t estimate(S) { + return 0; +} + +// Calculate the size of a SpecificStats's derived class. +// We need a template argument here rather than passing const SpecificStats& +// as we need to know the exact type to properly compute the size of the object. +template <typename S, std::enable_if_t<std::is_base_of_v<SpecificStats, S>, bool> = true> +inline size_t estimate(const S& stats) { + return stats.estimateObjectSizeInBytes() - sizeof(S); +} + +// Calculate the size of the inlined vector's elements. +template <typename T, size_t N, typename A> +size_t estimate(const absl::InlinedVector<T, N, A>& vector) { + size_t size = 0; + // Calculate size of the value only if the values are not inlined. + if (vector.capacity() > N) { + size += vector.capacity() * sizeof(T); + } + + for (const auto& elem : vector) { + size += estimate(elem); + } + + return size; +} + +// Calculate the size of the vector's elements. +template <typename T, typename A> +size_t estimate(const std::vector<T, A>& vector) { + size_t size = vector.capacity() * sizeof(T); + + for (const auto& elem : vector) { + size += estimate(elem); + } + + return size; +} + +template <typename T, typename A> +size_t estimateContainerOnly(const std::vector<T, A>& vector) { + return vector.capacity() * sizeof(T); +} + +template <typename K, typename V, typename... Args> +size_t estimate(const absl::flat_hash_map<K, V, Args...>& map) { + // The estimation is based on the memory usage of absl::flat_hash_map + // documented in https://abseil.io/docs/cpp/guides/container: + // The container uses O((sizeof(std::pair<const K, V>) + 1) * bucket_count()) bytes. + // The tests with a custom allocator showed that actual memory usage was + // (sizeof(std::pair<const K, V>) + 1) * bucket_count() + C, + // where C was equal to 17 for non-empty containers on x64 platform + // and 0 for empty containers. + constexpr size_t kEstimatedConstantPayload = 17; + size_t bucketSize = sizeof(std::pair<const K, V>) + 1; + size_t size = map.bucket_count() * bucketSize; + size += map.empty() ? 0 : kEstimatedConstantPayload; + + for (auto&& [key, val] : map) { + size += estimate(key); + size += estimate(val); + } + + return size; +} + +template <class BufferAllocator> +size_t estimate(const BasicBufBuilder<BufferAllocator>& ba) { + return static_cast<size_t>(ba.capacity()); +} +} // namespace mongo::sbe::size_estimator diff --git a/src/mongo/db/exec/sbe/stages/branch.cpp b/src/mongo/db/exec/sbe/stages/branch.cpp index 934e3095b9b..bec12b12ee2 100644 --- a/src/mongo/db/exec/sbe/stages/branch.cpp +++ b/src/mongo/db/exec/sbe/stages/branch.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/branch.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" namespace mongo { namespace sbe { @@ -242,5 +243,16 @@ std::vector<DebugPrinter::Block> BranchStage::debugPrint() const { return ret; } +size_t BranchStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += _filter->estimateSize(); + size += size_estimator::estimate(_inputThenVals); + size += size_estimator::estimate(_inputElseVals); + size += size_estimator::estimate(_outputVals); + size += size_estimator::estimate(_specificStats); + return size; +} + } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/branch.h b/src/mongo/db/exec/sbe/stages/branch.h index cababdef765..67b5af8a517 100644 --- a/src/mongo/db/exec/sbe/stages/branch.h +++ b/src/mongo/db/exec/sbe/stages/branch.h @@ -65,6 +65,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: const std::unique_ptr<EExpression> _filter; diff --git a/src/mongo/db/exec/sbe/stages/bson_scan.cpp b/src/mongo/db/exec/sbe/stages/bson_scan.cpp index cc360ccb4f6..c340071ba0e 100644 --- a/src/mongo/db/exec/sbe/stages/bson_scan.cpp +++ b/src/mongo/db/exec/sbe/stages/bson_scan.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/bson_scan.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/util/str.h" namespace mongo { @@ -181,5 +182,13 @@ std::vector<DebugPrinter::Block> BSONScanStage::debugPrint() const { return ret; } + +size_t BSONScanStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_fields); + size += size_estimator::estimate(_vars); + return size; +} + } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/bson_scan.h b/src/mongo/db/exec/sbe/stages/bson_scan.h index 14e650607d5..7804bcd4149 100644 --- a/src/mongo/db/exec/sbe/stages/bson_scan.h +++ b/src/mongo/db/exec/sbe/stages/bson_scan.h @@ -65,6 +65,7 @@ public: const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: const char* const _bsonBegin; diff --git a/src/mongo/db/exec/sbe/stages/check_bounds.cpp b/src/mongo/db/exec/sbe/stages/check_bounds.cpp index 2bc86c2ccd5..e5129b63fe1 100644 --- a/src/mongo/db/exec/sbe/stages/check_bounds.cpp +++ b/src/mongo/db/exec/sbe/stages/check_bounds.cpp @@ -31,6 +31,8 @@ #include "mongo/db/exec/sbe/stages/check_bounds.h" +#include "mongo/db/exec/sbe/size_estimator.h" + namespace mongo::sbe { CheckBoundsStage::CheckBoundsStage(std::unique_ptr<PlanStage> input, const CheckBoundsParams& params, @@ -174,6 +176,18 @@ std::vector<DebugPrinter::Block> CheckBoundsStage::debugPrint() const { return ret; } +size_t CheckBoundsStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_specificStats); + size += size_estimator::estimate(_params.keyPattern); + size += size_estimator::estimate(_params.bounds); + size += size_estimator::estimate(_seekPoint); + size += size_estimator::estimate(_checker); + size += size_estimator::estimate(_keyBuffer); + return size; +} + void CheckBoundsStage::doSaveState() { if (!slotsAccessible()) { return; diff --git a/src/mongo/db/exec/sbe/stages/check_bounds.h b/src/mongo/db/exec/sbe/stages/check_bounds.h index ae9a4bb3171..5201a41e2bb 100644 --- a/src/mongo/db/exec/sbe/stages/check_bounds.h +++ b/src/mongo/db/exec/sbe/stages/check_bounds.h @@ -83,6 +83,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; diff --git a/src/mongo/db/exec/sbe/stages/co_scan.cpp b/src/mongo/db/exec/sbe/stages/co_scan.cpp index 86e460ee208..73e89a5e87e 100644 --- a/src/mongo/db/exec/sbe/stages/co_scan.cpp +++ b/src/mongo/db/exec/sbe/stages/co_scan.cpp @@ -75,4 +75,8 @@ void CoScanStage::close() { trackClose(); } +size_t CoScanStage::estimateCompileTimeSize() const { + return sizeof(*this); +} + } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/co_scan.h b/src/mongo/db/exec/sbe/stages/co_scan.h index 74125ffe685..4625b636a14 100644 --- a/src/mongo/db/exec/sbe/stages/co_scan.h +++ b/src/mongo/db/exec/sbe/stages/co_scan.h @@ -54,5 +54,6 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; + size_t estimateCompileTimeSize() const final; }; } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/exchange.cpp b/src/mongo/db/exec/sbe/stages/exchange.cpp index baeba0fdc90..95b143ebec5 100644 --- a/src/mongo/db/exec/sbe/stages/exchange.cpp +++ b/src/mongo/db/exec/sbe/stages/exchange.cpp @@ -33,6 +33,7 @@ #include "mongo/base/init.h" #include "mongo/db/client.h" +#include "mongo/db/exec/sbe/size_estimator.h" namespace mongo::sbe { std::unique_ptr<ThreadPool> s_globalThreadPool; @@ -132,6 +133,16 @@ ExchangePipe* ExchangeState::pipe(size_t consumerTid, size_t producerTid) { return _consumers[consumerTid]->pipe(producerTid); } +size_t ExchangeState::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_fields); + size += _partition ? _partition->estimateSize() : 0; + size += _orderLess ? _orderLess->estimateSize() : 0; + size += size_estimator::estimate(_consumers); + size += size_estimator::estimate(_producers); + return size; +} + ExchangeBuffer* ExchangeConsumer::getBuffer(size_t producerId) { if (_fullBuffers[producerId]) { return _fullBuffers[producerId].get(); @@ -431,6 +442,13 @@ ExchangePipe* ExchangeConsumer::pipe(size_t producerTid) { } } +size_t ExchangeConsumer::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += _state->estimateCompileTimeSize(); + return size; +} + ExchangeBuffer* ExchangeProducer::getBuffer(size_t consumerId) { if (_emptyBuffers[consumerId]) { return _emptyBuffers[consumerId].get(); diff --git a/src/mongo/db/exec/sbe/stages/exchange.h b/src/mongo/db/exec/sbe/stages/exchange.h index a5f166c78a6..b94b4968f66 100644 --- a/src/mongo/db/exec/sbe/stages/exchange.h +++ b/src/mongo/db/exec/sbe/stages/exchange.h @@ -222,6 +222,8 @@ public: ExchangePipe* pipe(size_t consumerTid, size_t producerTid); + size_t estimateCompileTimeSize() const; + private: const ExchangePolicy _policy; const size_t _numOfProducers; @@ -276,6 +278,7 @@ public: std::vector<DebugPrinter::Block> debugPrint() const final; ExchangePipe* pipe(size_t producerTid); + size_t estimateCompileTimeSize() const final; private: ExchangeBuffer* getBuffer(size_t producerId); @@ -325,6 +328,12 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; + // This function should never be executed + // since ExchangeProducer is not created in compile time. + size_t estimateCompileTimeSize() const final { + MONGO_UNREACHABLE; + } + private: ExchangeBuffer* getBuffer(size_t consumerId); void putBuffer(size_t consumerId); diff --git a/src/mongo/db/exec/sbe/stages/filter.h b/src/mongo/db/exec/sbe/stages/filter.h index fabf41cf31a..2120be1c062 100644 --- a/src/mongo/db/exec/sbe/stages/filter.h +++ b/src/mongo/db/exec/sbe/stages/filter.h @@ -30,6 +30,7 @@ #pragma once #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/sbe/stages/stages.h" #include "mongo/db/exec/sbe/vm/vm.h" @@ -176,6 +177,14 @@ public: return ret; } + size_t estimateCompileTimeSize() const final { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += _filter->estimateSize(); + size += size_estimator::estimate(_specificStats); + return size; + } + private: const std::unique_ptr<EExpression> _filter; std::unique_ptr<vm::CodeFragment> _filterCode; diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp index e54d18e41bf..5ae9c234c81 100644 --- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp +++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp @@ -33,6 +33,8 @@ #include "mongo/util/str.h" +#include "mongo/db/exec/sbe/size_estimator.h" + namespace mongo { namespace sbe { HashAggStage::HashAggStage(std::unique_ptr<PlanStage> input, @@ -322,5 +324,15 @@ std::vector<DebugPrinter::Block> HashAggStage::debugPrint() const { return ret; } + +size_t HashAggStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_gbs); + size += size_estimator::estimate(_aggs); + size += size_estimator::estimate(_seekKeysSlots); + return size; +} + } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.h b/src/mongo/db/exec/sbe/stages/hash_agg.h index ff8120044f2..67302aeb02a 100644 --- a/src/mongo/db/exec/sbe/stages/hash_agg.h +++ b/src/mongo/db/exec/sbe/stages/hash_agg.h @@ -86,6 +86,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: using TableType = stdx::unordered_map<value::MaterializedRow, diff --git a/src/mongo/db/exec/sbe/stages/hash_join.cpp b/src/mongo/db/exec/sbe/stages/hash_join.cpp index 03f455414e6..86675029c0e 100644 --- a/src/mongo/db/exec/sbe/stages/hash_join.cpp +++ b/src/mongo/db/exec/sbe/stages/hash_join.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/hash_join.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/util/str.h" namespace mongo { @@ -290,5 +291,15 @@ std::vector<DebugPrinter::Block> HashJoinStage::debugPrint() const { return ret; } + +size_t HashJoinStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_outerCond); + size += size_estimator::estimate(_outerProjects); + size += size_estimator::estimate(_innerCond); + size += size_estimator::estimate(_innerProjects); + return size; +} } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/hash_join.h b/src/mongo/db/exec/sbe/stages/hash_join.h index f8ed199ff30..ed4781116d9 100644 --- a/src/mongo/db/exec/sbe/stages/hash_join.h +++ b/src/mongo/db/exec/sbe/stages/hash_join.h @@ -79,6 +79,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: using TableType = std::unordered_multimap<value::MaterializedRow, // NOLINT diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp index 76ce847de0a..cbe7e5867dd 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp +++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp @@ -33,6 +33,7 @@ #include "mongo/db/catalog/index_catalog.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/sbe/values/bson.h" #include "mongo/db/exec/trial_run_tracker.h" #include "mongo/db/index/index_access_method.h" @@ -494,4 +495,14 @@ std::vector<DebugPrinter::Block> IndexScanStage::debugPrint() const { return ret; } + +size_t IndexScanStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_vars); + size += size_estimator::estimate(_indexName); + size += size_estimator::estimate(_valuesBuffer); + size += size_estimator::estimate(_specificStats); + return size; +} + } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h index b542d225663..72c81bde98c 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.h +++ b/src/mongo/db/exec/sbe/stages/ix_scan.h @@ -96,6 +96,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() override; diff --git a/src/mongo/db/exec/sbe/stages/limit_skip.cpp b/src/mongo/db/exec/sbe/stages/limit_skip.cpp index 7e2206b8023..359355582ac 100644 --- a/src/mongo/db/exec/sbe/stages/limit_skip.cpp +++ b/src/mongo/db/exec/sbe/stages/limit_skip.cpp @@ -31,6 +31,8 @@ #include "mongo/db/exec/sbe/stages/limit_skip.h" +#include "mongo/db/exec/sbe/size_estimator.h" + namespace mongo::sbe { LimitSkipStage::LimitSkipStage(std::unique_ptr<PlanStage> input, boost::optional<long long> limit, @@ -127,4 +129,11 @@ std::vector<DebugPrinter::Block> LimitSkipStage::debugPrint() const { return ret; } + +size_t LimitSkipStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_specificStats); + return size; +} } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/limit_skip.h b/src/mongo/db/exec/sbe/stages/limit_skip.h index 127b509f58c..f0f62b34239 100644 --- a/src/mongo/db/exec/sbe/stages/limit_skip.h +++ b/src/mongo/db/exec/sbe/stages/limit_skip.h @@ -63,6 +63,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: const boost::optional<long long> _limit; diff --git a/src/mongo/db/exec/sbe/stages/loop_join.cpp b/src/mongo/db/exec/sbe/stages/loop_join.cpp index ce232437d72..54db378bdde 100644 --- a/src/mongo/db/exec/sbe/stages/loop_join.cpp +++ b/src/mongo/db/exec/sbe/stages/loop_join.cpp @@ -31,6 +31,7 @@ #include "mongo/db/exec/sbe/stages/loop_join.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/util/str.h" namespace mongo::sbe { @@ -241,4 +242,14 @@ std::vector<DebugPrinter::Block> LoopJoinStage::debugPrint() const { return ret; } + +size_t LoopJoinStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_outerProjects); + size += size_estimator::estimate(_outerCorrelated); + size += _predicate ? _predicate->estimateSize() : 0; + size += size_estimator::estimate(_specificStats); + return size; +} } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/loop_join.h b/src/mongo/db/exec/sbe/stages/loop_join.h index ffba27e9930..b335f880b24 100644 --- a/src/mongo/db/exec/sbe/stages/loop_join.h +++ b/src/mongo/db/exec/sbe/stages/loop_join.h @@ -74,6 +74,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: PlanState getNextOuterSide() { diff --git a/src/mongo/db/exec/sbe/stages/makeobj.cpp b/src/mongo/db/exec/sbe/stages/makeobj.cpp index 0281d835414..36c022544b1 100644 --- a/src/mongo/db/exec/sbe/stages/makeobj.cpp +++ b/src/mongo/db/exec/sbe/stages/makeobj.cpp @@ -31,6 +31,7 @@ #include "mongo/db/exec/sbe/stages/makeobj.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/sbe/values/bson.h" #include "mongo/util/str.h" @@ -409,6 +410,16 @@ std::vector<DebugPrinter::Block> MakeObjStageBase<O>::debugPrint() const { } template <MakeObjOutputType O> +size_t MakeObjStageBase<O>::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_fields); + size += size_estimator::estimate(_projectFields); + size += size_estimator::estimate(_projectVars); + return size; +} + +template <MakeObjOutputType O> void MakeObjStageBase<O>::doSaveState() { if (!slotsAccessible()) { return; diff --git a/src/mongo/db/exec/sbe/stages/makeobj.h b/src/mongo/db/exec/sbe/stages/makeobj.h index cc00916aa87..8fd81185e08 100644 --- a/src/mongo/db/exec/sbe/stages/makeobj.h +++ b/src/mongo/db/exec/sbe/stages/makeobj.h @@ -100,6 +100,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; diff --git a/src/mongo/db/exec/sbe/stages/merge_join.cpp b/src/mongo/db/exec/sbe/stages/merge_join.cpp index c2417022f79..37991cded4e 100644 --- a/src/mongo/db/exec/sbe/stages/merge_join.cpp +++ b/src/mongo/db/exec/sbe/stages/merge_join.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/merge_join.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/util/str.h" namespace mongo { @@ -423,5 +424,16 @@ std::vector<DebugPrinter::Block> MergeJoinStage::debugPrint() const { return ret; } + +size_t MergeJoinStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_outerKeys); + size += size_estimator::estimate(_outerProjects); + size += size_estimator::estimate(_innerKeys); + size += size_estimator::estimate(_innerProjects); + size += size_estimator::estimate(_dirs); + return size; +} } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/merge_join.h b/src/mongo/db/exec/sbe/stages/merge_join.h index 51edfc5b431..42425394bc0 100644 --- a/src/mongo/db/exec/sbe/stages/merge_join.h +++ b/src/mongo/db/exec/sbe/stages/merge_join.h @@ -75,6 +75,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; diff --git a/src/mongo/db/exec/sbe/stages/project.cpp b/src/mongo/db/exec/sbe/stages/project.cpp index d2a2ffb5fe1..ef6c6739ade 100644 --- a/src/mongo/db/exec/sbe/stages/project.cpp +++ b/src/mongo/db/exec/sbe/stages/project.cpp @@ -31,6 +31,8 @@ #include "mongo/db/exec/sbe/stages/project.h" +#include "mongo/db/exec/sbe/size_estimator.h" + namespace mongo { namespace sbe { ProjectStage::ProjectStage(std::unique_ptr<PlanStage> input, @@ -145,6 +147,13 @@ std::vector<DebugPrinter::Block> ProjectStage::debugPrint() const { return ret; } +size_t ProjectStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_projects); + return size; +} + void ProjectStage::doSaveState() { if (!slotsAccessible()) { return; diff --git a/src/mongo/db/exec/sbe/stages/project.h b/src/mongo/db/exec/sbe/stages/project.h index 9875d4c805a..24033b55aca 100644 --- a/src/mongo/db/exec/sbe/stages/project.h +++ b/src/mongo/db/exec/sbe/stages/project.h @@ -60,6 +60,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp index d236091af14..1e1cc4b3744 100644 --- a/src/mongo/db/exec/sbe/stages/scan.cpp +++ b/src/mongo/db/exec/sbe/stages/scan.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/scan.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/trial_run_tracker.h" #include "mongo/db/index/index_access_method.h" #include "mongo/db/repl/optime.h" @@ -483,6 +484,14 @@ std::vector<DebugPrinter::Block> ScanStage::debugPrint() const { return ret; } +size_t ScanStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_fields); + size += size_estimator::estimate(_vars); + size += size_estimator::estimate(_specificStats); + return size; +} + ParallelScanStage::ParallelScanStage(CollectionUUID collectionUuid, boost::optional<value::SlotId> recordSlot, boost::optional<value::SlotId> recordIdSlot, @@ -895,5 +904,13 @@ std::vector<DebugPrinter::Block> ParallelScanStage::debugPrint() const { return ret; } + +size_t ParallelScanStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_fields); + size += size_estimator::estimate(_vars); + return size; +} + } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/scan.h b/src/mongo/db/exec/sbe/stages/scan.h index 82e33c2334f..23d8370ce47 100644 --- a/src/mongo/db/exec/sbe/stages/scan.h +++ b/src/mongo/db/exec/sbe/stages/scan.h @@ -119,6 +119,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() override; @@ -228,6 +229,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; diff --git a/src/mongo/db/exec/sbe/stages/sort.cpp b/src/mongo/db/exec/sbe/stages/sort.cpp index 8f59e77baaa..a7476eafa0b 100644 --- a/src/mongo/db/exec/sbe/stages/sort.cpp +++ b/src/mongo/db/exec/sbe/stages/sort.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/sort.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/trial_run_tracker.h" #include "mongo/db/stats/resource_consumption_metrics.h" #include "mongo/util/str.h" @@ -305,5 +306,15 @@ std::vector<DebugPrinter::Block> SortStage::debugPrint() const { return ret; } + +size_t SortStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_obs); + size += size_estimator::estimate(_dirs); + size += size_estimator::estimate(_vals); + size += size_estimator::estimate(_specificStats); + return size; +} } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/sort.h b/src/mongo/db/exec/sbe/stages/sort.h index 05350171383..5599eedea64 100644 --- a/src/mongo/db/exec/sbe/stages/sort.h +++ b/src/mongo/db/exec/sbe/stages/sort.h @@ -85,6 +85,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doDetachFromTrialRunTracker() override; diff --git a/src/mongo/db/exec/sbe/stages/sorted_merge.cpp b/src/mongo/db/exec/sbe/stages/sorted_merge.cpp index 77e0eb65c23..f0a648f38ad 100644 --- a/src/mongo/db/exec/sbe/stages/sorted_merge.cpp +++ b/src/mongo/db/exec/sbe/stages/sorted_merge.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/sorted_merge.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" namespace mongo { namespace sbe { @@ -232,5 +233,15 @@ std::vector<DebugPrinter::Block> SortedMergeStage::debugPrint() const { return ret; } + +size_t SortedMergeStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_inputKeys); + size += size_estimator::estimate(_dirs); + size += size_estimator::estimate(_inputVals); + size += size_estimator::estimate(_outputVals); + return size; +} } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/sorted_merge.h b/src/mongo/db/exec/sbe/stages/sorted_merge.h index 8abeddad214..3b87e4c8849 100644 --- a/src/mongo/db/exec/sbe/stages/sorted_merge.h +++ b/src/mongo/db/exec/sbe/stages/sorted_merge.h @@ -74,6 +74,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: const std::vector<value::SlotVector> _inputKeys; diff --git a/src/mongo/db/exec/sbe/stages/spool.cpp b/src/mongo/db/exec/sbe/stages/spool.cpp index d87525d5c06..a7841ca6714 100644 --- a/src/mongo/db/exec/sbe/stages/spool.cpp +++ b/src/mongo/db/exec/sbe/stages/spool.cpp @@ -160,6 +160,13 @@ std::vector<DebugPrinter::Block> SpoolEagerProducerStage::debugPrint() const { return ret; } +size_t SpoolEagerProducerStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_vals); + return size; +} + SpoolLazyProducerStage::SpoolLazyProducerStage(std::unique_ptr<PlanStage> input, SpoolId spoolId, value::SlotVector vals, @@ -331,4 +338,12 @@ std::vector<DebugPrinter::Block> SpoolLazyProducerStage::debugPrint() const { DebugPrinter::addBlocks(ret, _children[0]->debugPrint()); return ret; } + +size_t SpoolLazyProducerStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_vals); + size += _predicate ? _predicate->estimateSize() : 0; + return size; +} } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/spool.h b/src/mongo/db/exec/sbe/stages/spool.h index 3daa00d1df4..197ad48d3ba 100644 --- a/src/mongo/db/exec/sbe/stages/spool.h +++ b/src/mongo/db/exec/sbe/stages/spool.h @@ -30,6 +30,7 @@ #pragma once #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/sbe/stages/stages.h" namespace mongo::sbe { @@ -68,6 +69,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: std::shared_ptr<SpoolBuffer> _buffer{nullptr}; @@ -120,6 +122,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; @@ -271,6 +274,12 @@ public: return ret; } + size_t estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_vals); + return size; + } + private: std::shared_ptr<SpoolBuffer> _buffer{nullptr}; size_t _bufferIt{0}; diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h index bf305b09513..5054618ca1f 100644 --- a/src/mongo/db/exec/sbe/stages/stages.h +++ b/src/mongo/db/exec/sbe/stages/stages.h @@ -413,6 +413,13 @@ public: return {DebugPrinter::Block(str)}; } + /** + * Estimates the compile-time size of the current plan stage and its children (SBE Plan + * subtree). The compile-time size is the size of the SBE subtree before it has been prepared or + * executed. + */ + virtual size_t estimateCompileTimeSize() const = 0; + friend class CanSwitchOperationContext<PlanStage>; friend class CanChangeState<PlanStage>; friend class CanTrackStats<PlanStage>; diff --git a/src/mongo/db/exec/sbe/stages/traverse.cpp b/src/mongo/db/exec/sbe/stages/traverse.cpp index 851ffffc977..ec92c0525b8 100644 --- a/src/mongo/db/exec/sbe/stages/traverse.cpp +++ b/src/mongo/db/exec/sbe/stages/traverse.cpp @@ -29,6 +29,7 @@ #include "mongo/platform/basic.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/db/exec/sbe/stages/traverse.h" namespace mongo::sbe { @@ -381,4 +382,14 @@ std::vector<DebugPrinter::Block> TraverseStage::debugPrint() const { return ret; } + +size_t TraverseStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_correlatedSlots); + size += _fold ? _fold->estimateSize() : 0; + size += _final ? _final->estimateSize() : 0; + size += size_estimator::estimate(_specificStats); + return size; +} } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/traverse.h b/src/mongo/db/exec/sbe/stages/traverse.h index ef29e7e24db..f59dc1a6763 100644 --- a/src/mongo/db/exec/sbe/stages/traverse.h +++ b/src/mongo/db/exec/sbe/stages/traverse.h @@ -87,6 +87,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; diff --git a/src/mongo/db/exec/sbe/stages/union.cpp b/src/mongo/db/exec/sbe/stages/union.cpp index cbf288c489f..a661e6c579f 100644 --- a/src/mongo/db/exec/sbe/stages/union.cpp +++ b/src/mongo/db/exec/sbe/stages/union.cpp @@ -32,6 +32,7 @@ #include "mongo/db/exec/sbe/stages/union.h" #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/size_estimator.h" namespace mongo::sbe { UnionStage::UnionStage(PlanStage::Vector inputStages, @@ -217,6 +218,14 @@ std::vector<DebugPrinter::Block> UnionStage::debugPrint() const { return ret; } +size_t UnionStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size += size_estimator::estimate(_inputVals); + size += size_estimator::estimate(_outputVals); + return size; +} + void UnionStage::clearBranches() { while (!_remainingBranchesToDrain.empty()) { auto& branch = _remainingBranchesToDrain.front(); diff --git a/src/mongo/db/exec/sbe/stages/union.h b/src/mongo/db/exec/sbe/stages/union.h index 3714f71a5e6..2ec0ec73df9 100644 --- a/src/mongo/db/exec/sbe/stages/union.h +++ b/src/mongo/db/exec/sbe/stages/union.h @@ -66,6 +66,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: struct UnionBranch { diff --git a/src/mongo/db/exec/sbe/stages/unique.cpp b/src/mongo/db/exec/sbe/stages/unique.cpp index d11d9201b52..355927ff912 100644 --- a/src/mongo/db/exec/sbe/stages/unique.cpp +++ b/src/mongo/db/exec/sbe/stages/unique.cpp @@ -31,6 +31,8 @@ #include "mongo/db/exec/sbe/stages/unique.h" +#include "mongo/db/exec/sbe/size_estimator.h" + namespace mongo { namespace sbe { UniqueStage::UniqueStage(std::unique_ptr<PlanStage> input, @@ -132,5 +134,14 @@ std::vector<DebugPrinter::Block> UniqueStage::debugPrint() const { return ret; } + +size_t UniqueStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + size_estimator::estimate(_keySlots); + size += size_estimator::estimate(_specificStats); + return size; +} + } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/stages/unique.h b/src/mongo/db/exec/sbe/stages/unique.h index d697026a72a..1165743a0cc 100644 --- a/src/mongo/db/exec/sbe/stages/unique.h +++ b/src/mongo/db/exec/sbe/stages/unique.h @@ -66,6 +66,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; private: const value::SlotVector _keySlots; diff --git a/src/mongo/db/exec/sbe/stages/unwind.cpp b/src/mongo/db/exec/sbe/stages/unwind.cpp index 6bcf87bd176..8afc781044f 100644 --- a/src/mongo/db/exec/sbe/stages/unwind.cpp +++ b/src/mongo/db/exec/sbe/stages/unwind.cpp @@ -31,6 +31,7 @@ #include "mongo/db/exec/sbe/stages/unwind.h" +#include "mongo/db/exec/sbe/size_estimator.h" #include "mongo/util/str.h" namespace mongo::sbe { @@ -220,4 +221,10 @@ void UnwindStage::doRestoreState() { _inArrayAccessor.refresh(); } + +size_t UnwindStage::estimateCompileTimeSize() const { + size_t size = sizeof(*this); + size += size_estimator::estimate(_children); + return size; +} } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/unwind.h b/src/mongo/db/exec/sbe/stages/unwind.h index c9afe10417a..149130cf415 100644 --- a/src/mongo/db/exec/sbe/stages/unwind.h +++ b/src/mongo/db/exec/sbe/stages/unwind.h @@ -65,6 +65,7 @@ public: std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final; const SpecificStats* getSpecificStats() const final; std::vector<DebugPrinter::Block> debugPrint() const final; + size_t estimateCompileTimeSize() const final; protected: void doSaveState() final; diff --git a/src/mongo/db/query/index_bounds.h b/src/mongo/db/query/index_bounds.h index 056dbe6e513..5bc2967bac4 100644 --- a/src/mongo/db/query/index_bounds.h +++ b/src/mongo/db/query/index_bounds.h @@ -176,6 +176,11 @@ struct IndexBounds { BoundInclusion boundInclusion; }; +class IndexBoundsChecker; +namespace sbe::size_estimator { +size_t estimate(const IndexBoundsChecker&); +} // namespace sbe::size_estimator + /** * A helper used by IndexScan to navigate an index. */ @@ -301,6 +306,8 @@ private: std::vector<int> _expectedDirection; std::vector<BSONElement> _keyValues; + + friend size_t sbe::size_estimator::estimate(const IndexBoundsChecker&); }; } // namespace mongo diff --git a/src/mongo/db/query/sbe_plan_cache.h b/src/mongo/db/query/sbe_plan_cache.h index cf49ca35ba7..eb546fbc006 100644 --- a/src/mongo/db/query/sbe_plan_cache.h +++ b/src/mongo/db/query/sbe_plan_cache.h @@ -89,8 +89,7 @@ struct CachedSbePlan { } uint64_t estimateObjectSizeInBytes() const { - // TODO SERVER-59331: handle size estimation. - return 0; + return root->estimateCompileTimeSize(); } std::unique_ptr<sbe::PlanStage> root; |