diff options
Diffstat (limited to 'src/mongo/db/exec/sbe')
35 files changed, 588 insertions, 1851 deletions
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript index f8246baf556..b99625b0833 100644 --- a/src/mongo/db/exec/sbe/SConscript +++ b/src/mongo/db/exec/sbe/SConscript @@ -23,7 +23,7 @@ env.Library( '$BUILD_DIR/mongo/base', '$BUILD_DIR/mongo/db/exec/js_function', '$BUILD_DIR/mongo/db/fts/base_fts', - '$BUILD_DIR/mongo/db/index/index_access_method', + '$BUILD_DIR/mongo/db/index/key_generator', '$BUILD_DIR/mongo/db/query/collation/collator_interface', '$BUILD_DIR/mongo/db/query/datetime/date_time_support', '$BUILD_DIR/mongo/db/query/query_index_bounds', @@ -59,9 +59,8 @@ sbeEnv.Library( LIBDEPS_PRIVATE=[ '$BUILD_DIR/mongo/db/bson/dotted_path_support', '$BUILD_DIR/mongo/db/sorter/sorter_idl', - '$BUILD_DIR/mongo/db/sorter/sorter_stats', - ], -) + ] + ) sbeEnv.Library( target='query_sbe_stages', @@ -90,6 +89,7 @@ sbeEnv.Library( LIBDEPS=[ '$BUILD_DIR/mongo/base', '$BUILD_DIR/mongo/db/concurrency/lock_manager', + '$BUILD_DIR/mongo/db/concurrency/write_conflict_exception', '$BUILD_DIR/mongo/db/exec/js_function', '$BUILD_DIR/mongo/db/exec/scoped_timer', '$BUILD_DIR/mongo/db/query/plan_yield_policy', @@ -106,7 +106,6 @@ sbeEnv.Library( LIBDEPS_PRIVATE=[ '$BUILD_DIR/mongo/db/bson/dotted_path_support', '$BUILD_DIR/mongo/db/sorter/sorter_idl', - '$BUILD_DIR/mongo/db/sorter/sorter_stats', 'query_sbe', 'query_sbe_storage', ] diff --git a/src/mongo/db/exec/sbe/abt/abt_lower.cpp b/src/mongo/db/exec/sbe/abt/abt_lower.cpp index 7939d166a78..7da6a8e2cef 100644 --- a/src/mongo/db/exec/sbe/abt/abt_lower.cpp +++ b/src/mongo/db/exec/sbe/abt/abt_lower.cpp @@ -586,15 +586,14 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const GroupByNode& n, auto& names = binderAgg->names(); auto& exprs = refsAgg->nodes(); - sbe::SlotExprPairVector aggs; - aggs.reserve(exprs.size()); + sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> aggs; for (size_t idx = 0; idx < exprs.size(); ++idx) { auto expr = SBEExpressionLowering{_env, _slotMap}.optimize(exprs[idx]); auto slot = _slotIdGenerator.generate(); _slotMap.emplace(names[idx], slot); - aggs.push_back({slot, std::move(expr)}); + aggs.emplace(slot, std::move(expr)); } // TODO: use collator slot. @@ -610,11 +609,6 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const GroupByNode& n, true /*optimizedClose*/, collatorSlot, false /*allowDiskUse*/, - // Since we are always disallowing disk use for this stage, - // we need not provide merging expressions. Once spilling - // is permitted here, we will need to generate merging - // expressions during lowering. - sbe::makeSlotExprPairVec() /*mergingExprs*/, planNodeId); } @@ -1011,7 +1005,6 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const IndexScanNode& n, co resultSlot, ridSlot, boost::none, - boost::none, indexKeysToInclude, vars, seekKeySlotLower, diff --git a/src/mongo/db/exec/sbe/expression_test_base.h b/src/mongo/db/exec/sbe/expression_test_base.h index bc5121407f9..72d9820d760 100644 --- a/src/mongo/db/exec/sbe/expression_test_base.h +++ b/src/mongo/db/exec/sbe/expression_test_base.h @@ -69,24 +69,6 @@ protected: } /** - * Compiles 'expr' to bytecode when 'expr' is computing an aggregate. The current aggregate - * value can be read out of the provided 'aggAccessor'. - * - * Note that when actually executing the resulting bytecode, the caller is responsible for - * setting the value of 'aggAccessor' to the new resulting aggregate value. - */ - std::unique_ptr<vm::CodeFragment> compileAggExpression(const EExpression& expr, - value::SlotAccessor* aggAccessor) { - ON_BLOCK_EXIT([this] { - _ctx.aggExpression = false; - _ctx.accumulator = nullptr; - }); - _ctx.aggExpression = true; - _ctx.accumulator = aggAccessor; - return expr.compile(_ctx); - } - - /** * The caller takes ownership of the Value returned by this function and must call * 'releaseValue()' on it. The preferred way to ensure the Value is properly released is to * immediately store it in a ValueGuard. diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp index 65421e8b373..3a2b20a4657 100644 --- a/src/mongo/db/exec/sbe/expressions/expression.cpp +++ b/src/mongo/db/exec/sbe/expressions/expression.cpp @@ -429,8 +429,6 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = { BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::doubleDoubleSum, false}}, {"aggDoubleDoubleSum", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggDoubleDoubleSum, true}}, - {"aggMergeDoubleDoubleSums", - BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggMergeDoubleDoubleSums, true}}, {"doubleDoubleSumFinalize", BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::doubleDoubleSumFinalize, false}}, {"doubleDoubleMergeSumFinalize", @@ -438,8 +436,6 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = { {"doubleDoublePartialSumFinalize", BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::doubleDoublePartialSumFinalize, false}}, {"aggStdDev", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggStdDev, true}}, - {"aggMergeStdDevs", - BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggMergeStdDevs, true}}, {"stdDevPopFinalize", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::stdDevPopFinalize, false}}, {"stdDevSampFinalize", @@ -471,10 +467,7 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = { {"tan", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::tan, false}}, {"tanh", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::tanh, false}}, {"round", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::round, false}}, - {"concat", BuiltinFn{kAnyNumberOfArgs, vm::Builtin::concat, false}}, - {"concatArrays", BuiltinFn{kAnyNumberOfArgs, vm::Builtin::concatArrays, false}}, - {"aggConcatArraysCapped", - BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::aggConcatArraysCapped, true}}, + {"concat", BuiltinFn{[](size_t n) { return n > 0; }, vm::Builtin::concat, false}}, {"isMember", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::isMember, false}}, {"collIsMember", BuiltinFn{[](size_t n) { return n == 3; }, vm::Builtin::collIsMember, false}}, {"indexOfBytes", @@ -493,11 +486,6 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = { BuiltinFn{[](size_t n) { return n >= 1; }, vm::Builtin::collSetIntersection, false}}, {"collSetDifference", BuiltinFn{[](size_t n) { return n == 3; }, vm::Builtin::collSetDifference, false}}, - {"aggSetUnion", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggSetUnion, true}}, - {"aggSetUnionCapped", - BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::aggSetUnionCapped, true}}, - {"aggCollSetUnionCapped", - BuiltinFn{[](size_t n) { return n == 3; }, vm::Builtin::aggCollSetUnionCapped, true}}, {"runJsPredicate", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::runJsPredicate, false}}, {"regexCompile", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::regexCompile, false}}, diff --git a/src/mongo/db/exec/sbe/expressions/expression.h b/src/mongo/db/exec/sbe/expressions/expression.h index d5a2c0fdf0e..2b9032f257b 100644 --- a/src/mongo/db/exec/sbe/expressions/expression.h +++ b/src/mongo/db/exec/sbe/expressions/expression.h @@ -348,8 +348,6 @@ private: std::string toString() const; }; -using SlotExprPairVector = std::vector<std::pair<value::SlotId, std::unique_ptr<EExpression>>>; - template <typename T, typename... Args> inline std::unique_ptr<EExpression> makeE(Args&&... args) { return std::make_unique<T>(std::forward<Args>(args)...); @@ -366,30 +364,20 @@ inline auto makeEs(Ts&&... pack) { namespace detail { // base case -template <typename R> -inline void makeSlotExprPairHelper(R& result, - value::SlotId slot, - std::unique_ptr<EExpression> expr) { - if constexpr (std::is_same_v<R, value::SlotMap<std::unique_ptr<EExpression>>>) { - result.emplace(slot, std::move(expr)); - } else { - result.push_back({slot, std::move(expr)}); - } +inline void makeEM_unwind(value::SlotMap<std::unique_ptr<EExpression>>& result, + value::SlotId slot, + std::unique_ptr<EExpression> expr) { + result.emplace(slot, std::move(expr)); } // recursive case -template <typename R, typename... Ts> -inline void makeSlotExprPairHelper(R& result, - value::SlotId slot, - std::unique_ptr<EExpression> expr, - Ts&&... rest) { - if constexpr (std::is_same_v<R, value::SlotMap<std::unique_ptr<EExpression>>>) { - result.emplace(slot, std::move(expr)); - } else { - static_assert(std::is_same_v<R, SlotExprPairVector>); - result.push_back({slot, std::move(expr)}); - } - makeSlotExprPairHelper(result, std::forward<Ts>(rest)...); +template <typename... Ts> +inline void makeEM_unwind(value::SlotMap<std::unique_ptr<EExpression>>& result, + value::SlotId slot, + std::unique_ptr<EExpression> expr, + Ts&&... rest) { + result.emplace(slot, std::move(expr)); + makeEM_unwind(result, std::forward<Ts>(rest)...); } } // namespace detail @@ -398,7 +386,7 @@ auto makeEM(Ts&&... pack) { value::SlotMap<std::unique_ptr<EExpression>> result; if constexpr (sizeof...(pack) > 0) { result.reserve(sizeof...(Ts) / 2); - detail::makeSlotExprPairHelper(result, std::forward<Ts>(pack)...); + detail::makeEM_unwind(result, std::forward<Ts>(pack)...); } return result; } @@ -411,16 +399,6 @@ auto makeSV(Args&&... args) { return v; } -template <typename... Ts> -auto makeSlotExprPairVec(Ts&&... pack) { - SlotExprPairVector v; - if constexpr (sizeof...(pack) > 0) { - v.reserve(sizeof...(Ts) / 2); - detail::makeSlotExprPairHelper(v, std::forward<Ts>(pack)...); - } - return v; -} - /** * This is a constant expression. It assumes the ownership of the input constant. */ diff --git a/src/mongo/db/exec/sbe/expressions/sbe_set_expressions_test.cpp b/src/mongo/db/exec/sbe/expressions/sbe_set_expressions_test.cpp index 279f7a287b4..30eb61a7b2d 100644 --- a/src/mongo/db/exec/sbe/expressions/sbe_set_expressions_test.cpp +++ b/src/mongo/db/exec/sbe/expressions/sbe_set_expressions_test.cpp @@ -28,7 +28,6 @@ */ #include "mongo/db/exec/sbe/expression_test_base.h" -#include "mongo/db/query/sbe_stage_builder_helpers.h" namespace mongo::sbe { @@ -94,35 +93,6 @@ TEST_F(SBEBuiltinSetOpTest, ReturnsNothingSetUnion) { runAndAssertNothing(compiledExpr.get()); } -TEST_F(SBEBuiltinSetOpTest, AggSetUnion) { - value::OwnedValueAccessor aggAccessor, inputAccessor; - auto inputSlot = bindAccessor(&inputAccessor); - auto setUnionExpr = - stage_builder::makeFunction("aggSetUnion", stage_builder::makeVariable(inputSlot)); - auto compiledExpr = compileAggExpression(*setUnionExpr, &aggAccessor); - - auto [arrTag1, arrVal1] = makeArray(BSON_ARRAY(1 << 2)); - inputAccessor.reset(arrTag1, arrVal1); - auto [resTag1, resVal1] = makeArraySet(BSON_ARRAY(1 << 2)); - runAndAssertExpression(compiledExpr.get(), {resTag1, resVal1}); - aggAccessor.reset(resTag1, resVal1); - - auto [arrTag2, arrVal2] = makeArraySet(BSON_ARRAY(1 << 3 << 2 << 6)); - inputAccessor.reset(arrTag2, arrVal2); - auto [resTag2, resVal2] = makeArraySet(BSON_ARRAY(1 << 2 << 3 << 6)); - runAndAssertExpression(compiledExpr.get(), {resTag2, resVal2}); - aggAccessor.reset(resTag2, resVal2); - - auto [arrTag3, arrVal3] = makeArray(BSONArray{}); - inputAccessor.reset(arrTag3, arrVal3); - auto [resTag3, resVal3] = makeArraySet(BSON_ARRAY(1 << 2 << 3 << 6)); - runAndAssertExpression(compiledExpr.get(), {resTag3, resVal3}); - aggAccessor.reset(resTag3, resVal3); - - inputAccessor.reset(value::TypeTags::Nothing, 0); - runAndAssertNothing(compiledExpr.get()); -} - TEST_F(SBEBuiltinSetOpTest, ComputesSetIntersection) { value::OwnedValueAccessor slotAccessor1, slotAccessor2; auto arrSlot1 = bindAccessor(&slotAccessor1); diff --git a/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp b/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp index 89244fd4589..87fc8dbd1f2 100644 --- a/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp +++ b/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp @@ -69,23 +69,18 @@ void HashAggStageTest::performHashAggWithSpillChecking( auto makeStageFn = [this, collatorSlot, shouldUseCollator, shouldSpill]( value::SlotId scanSlot, std::unique_ptr<PlanStage> scanStage) { auto countsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction("sum", + makeE<EConstant>(value::TypeTags::NumberInt64, + value::bitcastFrom<int64_t>(1)))), makeSV(), true, boost::optional<value::SlotId>{shouldUseCollator, collatorSlot}, shouldSpill, - makeSlotExprPairVec( - spillSlot, - stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); return std::make_pair(countsSlot, std::move(hashAggStage)); @@ -171,21 +166,20 @@ TEST_F(HashAggStageTest, HashAggMinMaxTest) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(), - makeSlotExprPairVec(minSlot, - stage_builder::makeFunction("min", makeE<EVariable>(scanSlot)), - maxSlot, - stage_builder::makeFunction("max", makeE<EVariable>(scanSlot)), - collMinSlot, - stage_builder::makeFunction( - "collMin", collExpr->clone(), makeE<EVariable>(scanSlot)), - collMaxSlot, - stage_builder::makeFunction( - "collMax", collExpr->clone(), makeE<EVariable>(scanSlot))), + makeEM(minSlot, + stage_builder::makeFunction("min", makeE<EVariable>(scanSlot)), + maxSlot, + stage_builder::makeFunction("max", makeE<EVariable>(scanSlot)), + collMinSlot, + stage_builder::makeFunction( + "collMin", collExpr->clone(), makeE<EVariable>(scanSlot)), + collMaxSlot, + stage_builder::makeFunction( + "collMax", collExpr->clone(), makeE<EVariable>(scanSlot))), makeSV(), true, boost::none, false /* allowDiskUse */, - makeSlotExprPairVec() /* mergingExprs */, kEmptyPlanNodeId); auto outSlot = generateSlotId(); @@ -236,15 +230,13 @@ TEST_F(HashAggStageTest, HashAggAddToSetTest) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(), - makeSlotExprPairVec(hashAggSlot, - stage_builder::makeFunction("collAddToSet", - std::move(collExpr), - makeE<EVariable>(scanSlot))), + makeEM(hashAggSlot, + stage_builder::makeFunction( + "collAddToSet", std::move(collExpr), makeE<EVariable>(scanSlot))), makeSV(), true, boost::none, false /* allowDiskUse */, - makeSlotExprPairVec() /* mergingExprs */, kEmptyPlanNodeId); return std::make_pair(hashAggSlot, std::move(hashAggStage)); @@ -335,16 +327,14 @@ TEST_F(HashAggStageTest, HashAggSeekKeysTest) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction("sum", + makeE<EConstant>(value::TypeTags::NumberInt64, + value::bitcastFrom<int64_t>(1)))), makeSV(seekSlot), true, boost::none, false /* allowDiskUse */, - makeSlotExprPairVec() /* mergingExprs */, kEmptyPlanNodeId); return std::make_pair(countsSlot, std::move(hashAggStage)); @@ -397,21 +387,17 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpill) { // Build a HashAggStage, group by the scanSlot and compute a simple count. auto countsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), makeSV(), // Seek slot true, boost::none, true /* allowDiskUse */, - makeSlotExprPairVec( - spillSlot, stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -434,7 +420,6 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpill) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_FALSE(stats->usedDisk); - ASSERT_EQ(0, stats->numSpills); ASSERT_EQ(0, stats->spilledRecords); stage->close(); @@ -443,6 +428,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpill) { TEST_F(HashAggStageTest, HashAggBasicCountSpill) { // We estimate the size of result row like {int64, int64} at 50B. Set the memory threshold to // 64B so that exactly one row fits in memory. + const int expectedRowsToFitInMemory = 1; auto defaultInternalQuerySBEAggApproxMemoryUseInBytesBeforeSpill = internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.load(); internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.store(64); @@ -460,21 +446,17 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpill) { // Build a HashAggStage, group by the scanSlot and compute a simple count. auto countsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), makeSV(), // Seek slot true, boost::none, true /* allowDiskUse */, - makeSlotExprPairVec( - spillSlot, stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -497,14 +479,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpill) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_TRUE(stats->usedDisk); - // Memory usage is estimated only every two rows at the most frequent. Also, we only start - // spilling after estimating that the memory budget is exceeded. These two factors result in - // fewer expected spills than there are input records, even though only one record fits in - // memory at a time. - ASSERT_EQ(stats->numSpills, 3); - // The input has one run of two consecutive values, so we expect to spill as many records as - // there are input values minus one. - ASSERT_EQ(stats->spilledRecords, 8); + ASSERT_EQ(results.size() - expectedRowsToFitInMemory, stats->spilledRecords); stage->close(); } @@ -538,21 +513,17 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillIfNoMemCheck) { // Build a HashAggStage, group by the scanSlot and compute a simple count. auto countsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), makeSV(), // Seek slot true, boost::none, true /* allowDiskUse */, - makeSlotExprPairVec( - spillSlot, stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -575,7 +546,6 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillIfNoMemCheck) { // Check that it did not spill. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_FALSE(stats->usedDisk); - ASSERT_EQ(0, stats->numSpills); ASSERT_EQ(0, stats->spilledRecords); stage->close(); @@ -584,6 +554,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillIfNoMemCheck) { TEST_F(HashAggStageTest, HashAggBasicCountSpillDouble) { // We estimate the size of result row like {double, int64} at 50B. Set the memory threshold to // 64B so that exactly one row fits in memory. + const int expectedRowsToFitInMemory = 1; auto defaultInternalQuerySBEAggApproxMemoryUseInBytesBeforeSpill = internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.load(); internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.store(64); @@ -601,21 +572,17 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpillDouble) { // Build a HashAggStage, group by the scanSlot and compute a simple count. auto countsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), makeSV(), // Seek slot true, boost::none, true /* allowDiskUse */, - makeSlotExprPairVec( - spillSlot, stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -638,14 +605,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpillDouble) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_TRUE(stats->usedDisk); - // Memory usage is estimated only every two rows at the most frequent. Also, we only start - // spilling after estimating that the memory budget is exceeded. These two factors result in - // fewer expected spills than there are input records, even though only one record fits in - // memory at a time. - ASSERT_EQ(stats->numSpills, 3); - // The input has one run of two consecutive values, so we expect to spill as many records as - // there are input values minus one. - ASSERT_EQ(stats->spilledRecords, 8); + ASSERT_EQ(results.size() - expectedRowsToFitInMemory, stats->spilledRecords); stage->close(); } @@ -667,21 +627,17 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillWithNoGroupByDouble) { // Build a HashAggStage, with an empty group by slot and compute a simple count. auto countsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), makeSV(), // Seek slot true, boost::none, true /* allowDiskUse */, - makeSlotExprPairVec( - spillSlot, stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -702,7 +658,6 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillWithNoGroupByDouble) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_FALSE(stats->usedDisk); - ASSERT_EQ(0, stats->numSpills); ASSERT_EQ(0, stats->spilledRecords); stage->close(); @@ -711,6 +666,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillWithNoGroupByDouble) { TEST_F(HashAggStageTest, HashAggMultipleAccSpill) { // We estimate the size of result row like {double, int64} at 59B. Set the memory threshold to // 128B so that two rows fit in memory. + const int expectedRowsToFitInMemory = 2; auto defaultInternalQuerySBEAggApproxMemoryUseInBytesBeforeSpill = internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.load(); internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.store(128); @@ -729,27 +685,19 @@ TEST_F(HashAggStageTest, HashAggMultipleAccSpill) { // Build a HashAggStage, group by the scanSlot and compute a simple count. auto countsSlot = generateSlotId(); auto sumsSlot = generateSlotId(); - auto spillSlot1 = generateSlotId(); - auto spillSlot2 = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1))), - sumsSlot, - stage_builder::makeFunction("sum", makeE<EVariable>(scanSlot))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1))), + sumsSlot, + stage_builder::makeFunction("sum", makeE<EVariable>(scanSlot))), makeSV(), // Seek slot true, boost::none, true /* allowDiskUse */, - makeSlotExprPairVec( - spillSlot1, - stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot1)), - spillSlot2, - stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot2))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -779,10 +727,7 @@ TEST_F(HashAggStageTest, HashAggMultipleAccSpill) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_TRUE(stats->usedDisk); - ASSERT_EQ(stats->numSpills, 3); - // The input has one run of two consecutive values, so we expect to spill as many records as - // there are input values minus one. - ASSERT_EQ(stats->spilledRecords, 8); + ASSERT_EQ(results.size() - expectedRowsToFitInMemory, stats->spilledRecords); stage->close(); } @@ -807,27 +752,19 @@ TEST_F(HashAggStageTest, HashAggMultipleAccSpillAllToDisk) { // Build a HashAggStage, group by the scanSlot and compute a simple count. auto countsSlot = generateSlotId(); auto sumsSlot = generateSlotId(); - auto spillSlot1 = generateSlotId(); - auto spillSlot2 = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1))), - sumsSlot, - stage_builder::makeFunction("sum", makeE<EVariable>(scanSlot))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1))), + sumsSlot, + stage_builder::makeFunction("sum", makeE<EVariable>(scanSlot))), makeSV(), // Seek slot true, boost::none, true, // allowDiskUse=true - makeSlotExprPairVec( - spillSlot1, - stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot1)), - spillSlot2, - stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot2))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -857,9 +794,7 @@ TEST_F(HashAggStageTest, HashAggMultipleAccSpillAllToDisk) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_TRUE(stats->usedDisk); - // We expect each incoming value to result in a spill of a single record. - ASSERT_EQ(stats->numSpills, 9); - ASSERT_EQ(stats->spilledRecords, 9); + ASSERT_EQ(results.size(), stats->spilledRecords); stage->close(); } @@ -896,18 +831,14 @@ TEST_F(HashAggStageTest, HashAggSum10Groups) { // Build a HashAggStage, group by the scanSlot and compute a sum for each group. auto sumsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec(sumsSlot, - stage_builder::makeFunction("sum", makeE<EVariable>(scanSlot))), + makeEM(sumsSlot, stage_builder::makeFunction("sum", makeE<EVariable>(scanSlot))), makeSV(), // Seek slot true, boost::none, true, // allowDiskUse=true - makeSlotExprPairVec( - spillSlot, stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. @@ -941,21 +872,17 @@ TEST_F(HashAggStageTest, HashAggBasicCountWithRecordIds) { // Build a HashAggStage, group by the scanSlot and compute a simple count. auto countsSlot = generateSlotId(); - auto spillSlot = generateSlotId(); auto stage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), makeSV(), // Seek slot true, boost::none, true, // allowDiskUse=true - makeSlotExprPairVec( - spillSlot, stage_builder::makeFunction("sum", stage_builder::makeVariable(spillSlot))), kEmptyPlanNodeId); // Prepare the tree and get the 'SlotAccessor' for the output slot. diff --git a/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp b/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp index f62b8a6229c..191f0ffe562 100644 --- a/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp +++ b/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp @@ -139,12 +139,11 @@ TEST_F(PlanSizeTest, Filter) { TEST_F(PlanSizeTest, HashAgg) { auto stage = makeS<HashAggStage>(mockS(), mockSV(), - makeSlotExprPairVec(generateSlotId(), mockE()), + makeEM(generateSlotId(), mockE()), makeSV(), true, generateSlotId(), false, - makeSlotExprPairVec(), kEmptyPlanNodeId); assertPlanSize(*stage); } @@ -169,7 +168,6 @@ TEST_F(PlanSizeTest, IndexScan) { generateSlotId(), generateSlotId(), generateSlotId(), - generateSlotId(), IndexKeysInclusionSet(1), mockSV(), generateSlotId(), diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp index 4f3e6d44372..77b1997bdee 100644 --- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp +++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp @@ -66,12 +66,7 @@ PlanStageTestFixture::generateVirtualScanMulti(int32_t numSlots, const BSONArray } void PlanStageTestFixture::prepareTree(CompileCtx* ctx, PlanStage* root) { - // We want to avoid recursive locking since this results in yield plans that don't yield when - // they should. - boost::optional<Lock::GlobalLock> globalLock; - if (!opCtx()->lockState()->isLocked()) { - globalLock.emplace(opCtx(), MODE_IS); - } + Lock::GlobalLock globalLock{opCtx(), MODE_IS}; root->attachToOperationContext(opCtx()); root->prepare(*ctx); root->open(false); diff --git a/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp b/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp index 57d3ff4e05f..54fadbe1f15 100644 --- a/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp +++ b/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp @@ -141,16 +141,14 @@ TEST_F(TrialRunTrackerTest, TrialEndsDuringOpenPhaseOfBlockingStage) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), - makeSV(), /* Seek slot */ + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSV(), // Seek slot true, boost::none, false /* allowDiskUse */, - makeSlotExprPairVec(), /* mergingExprs */ kEmptyPlanNodeId); auto tracker = std::make_unique<TrialRunTracker>(numResultsLimit, size_t{0}); @@ -212,16 +210,14 @@ TEST_F(TrialRunTrackerTest, OnlyDeepestNestedBlockingStageHasTrialRunTracker) { auto hashAggStage = makeS<HashAggStage>( std::move(unionStage), makeSV(unionSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), - makeSV(), /* Seek slot */ + makeEM(countsSlot, + stage_builder::makeFunction( + "sum", + makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSV(), // Seek slot true, boost::none, false /* allowDiskUse */, - makeSlotExprPairVec(), /* mergingExprs */ kEmptyPlanNodeId); hashAggStage->prepare(*ctx); @@ -281,16 +277,14 @@ TEST_F(TrialRunTrackerTest, SiblingBlockingStagesBothGetTrialRunTracker) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), - makeSV(), /* Seek slot */ + makeEM(countsSlot, + stage_builder::makeFunction("sum", + makeE<EConstant>(value::TypeTags::NumberInt64, + value::bitcastFrom<int64_t>(1)))), + makeSV(), // Seek slot true, boost::none, false /* allowDiskUse */, - makeSlotExprPairVec(), /* mergingExprs */ kEmptyPlanNodeId); return std::make_pair(countsSlot, std::move(hashAggStage)); @@ -413,16 +407,14 @@ TEST_F(TrialRunTrackerTest, DisablingTrackingForAChildStagePreventsEarlyExit) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeSlotExprPairVec( - countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), - makeSV(), /* Seek slot */ + makeEM(countsSlot, + stage_builder::makeFunction("sum", + makeE<EConstant>(value::TypeTags::NumberInt64, + value::bitcastFrom<int64_t>(1)))), + makeSV(), // Seek slot true, boost::none, false /* allowDiskUse */, - makeSlotExprPairVec(), /* mergingExprs */ kEmptyPlanNodeId); return std::make_pair(countsSlot, std::move(hashAggStage)); diff --git a/src/mongo/db/exec/sbe/size_estimator.h b/src/mongo/db/exec/sbe/size_estimator.h index bbb92328331..fb6684eea52 100644 --- a/src/mongo/db/exec/sbe/size_estimator.h +++ b/src/mongo/db/exec/sbe/size_estimator.h @@ -92,11 +92,6 @@ inline size_t estimate(const S& stats) { return stats.estimateObjectSizeInBytes() - sizeof(S); } -template <typename A, typename B> -inline size_t estimate(const std::pair<A, B>& pair) { - return estimate(pair.first) + estimate(pair.second); -} - // 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) { diff --git a/src/mongo/db/exec/sbe/stages/branch.cpp b/src/mongo/db/exec/sbe/stages/branch.cpp index 1e8809f0858..bec12b12ee2 100644 --- a/src/mongo/db/exec/sbe/stages/branch.cpp +++ b/src/mongo/db/exec/sbe/stages/branch.cpp @@ -70,29 +70,31 @@ void BranchStage::prepare(CompileCtx& ctx) { _children[0]->prepare(ctx); _children[1]->prepare(ctx); - // All of the slots listed in '_outputVals' must be unique. for (size_t idx = 0; idx < _outputVals.size(); ++idx) { - auto slot = _outputVals[idx]; - auto [_, inserted] = dupCheck.insert(slot); - uassert(4822831, str::stream() << "duplicate field: " << slot, inserted); - } + std::vector<value::SlotAccessor*> accessors; + accessors.reserve(2); - for (size_t idx = 0; idx < _outputVals.size(); ++idx) { - auto thenSlot = _inputThenVals[idx]; - auto elseSlot = _inputElseVals[idx]; + { + auto slot = _inputThenVals[idx]; + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822829, str::stream() << "duplicate field: " << slot, inserted); - // Slots listed in '_inputThenVals' and '_inputElseVals' may not appear in '_outputVals'. - bool thenSlotFound = dupCheck.count(thenSlot); - bool elseSlotFound = dupCheck.count(elseSlot); - uassert(4822829, str::stream() << "duplicate field: " << thenSlot, !thenSlotFound); - uassert(4822830, str::stream() << "duplicate field: " << elseSlot, !elseSlotFound); + accessors.emplace_back(_children[0]->getAccessor(ctx, slot)); + } + { + auto slot = _inputElseVals[idx]; + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822830, str::stream() << "duplicate field: " << slot, inserted); - std::vector<value::SlotAccessor*> accessors; - accessors.reserve(2); - accessors.emplace_back(_children[0]->getAccessor(ctx, thenSlot)); - accessors.emplace_back(_children[1]->getAccessor(ctx, elseSlot)); + accessors.emplace_back(_children[1]->getAccessor(ctx, slot)); + } + { + auto slot = _outputVals[idx]; + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822831, str::stream() << "duplicate field: " << slot, inserted); - _outValueAccessors.emplace_back(value::SwitchAccessor{std::move(accessors)}); + _outValueAccessors.emplace_back(value::SwitchAccessor{std::move(accessors)}); + } } // compile filter diff --git a/src/mongo/db/exec/sbe/stages/collection_helpers.h b/src/mongo/db/exec/sbe/stages/collection_helpers.h index 7471763a44a..4116f2980ff 100644 --- a/src/mongo/db/exec/sbe/stages/collection_helpers.h +++ b/src/mongo/db/exec/sbe/stages/collection_helpers.h @@ -40,7 +40,6 @@ namespace mongo::sbe { * A callback which gets called whenever a SCAN stage asks an underlying index scan for a result. */ using IndexKeyConsistencyCheckCallback = std::function<bool(OperationContext* opCtx, - StringMap<const IndexCatalogEntry*>&, value::SlotAccessor* snapshotIdAccessor, value::SlotAccessor* indexIdAccessor, value::SlotAccessor* indexKeyAccessor, diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp index 2514dedf1b0..14f40afeaa9 100644 --- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp +++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp @@ -27,67 +27,47 @@ * it in the license file. */ -#include "mongo/db/exec/sbe/stages/hash_agg.h" +#include "mongo/platform/basic.h" #include "mongo/db/concurrency/d_concurrency.h" -#include "mongo/db/exec/sbe/size_estimator.h" +#include "mongo/db/concurrency/write_conflict_exception.h" +#include "mongo/db/exec/sbe/stages/hash_agg.h" #include "mongo/db/exec/sbe/util/spilling.h" -#include "mongo/db/stats/resource_consumption_metrics.h" #include "mongo/db/storage/kv/kv_engine.h" #include "mongo/db/storage/storage_engine.h" #include "mongo/util/str.h" +#include "mongo/db/exec/sbe/size_estimator.h" + namespace mongo { namespace sbe { HashAggStage::HashAggStage(std::unique_ptr<PlanStage> input, value::SlotVector gbs, - SlotExprPairVector aggs, + value::SlotMap<std::unique_ptr<EExpression>> aggs, value::SlotVector seekKeysSlots, bool optimizedClose, boost::optional<value::SlotId> collatorSlot, bool allowDiskUse, - SlotExprPairVector mergingExprs, - PlanNodeId planNodeId, - bool forceIncreasedSpilling) + PlanNodeId planNodeId) : PlanStage("group"_sd, planNodeId), _gbs(std::move(gbs)), _aggs(std::move(aggs)), _collatorSlot(collatorSlot), _allowDiskUse(allowDiskUse), _seekKeysSlots(std::move(seekKeysSlots)), - _optimizedClose(optimizedClose), - _mergingExprs(std::move(mergingExprs)), - _forceIncreasedSpilling(forceIncreasedSpilling) { + _optimizedClose(optimizedClose) { _children.emplace_back(std::move(input)); invariant(_seekKeysSlots.empty() || _seekKeysSlots.size() == _gbs.size()); tassert(5843100, "HashAgg stage was given optimizedClose=false and seek keys", _seekKeysSlots.empty() || _optimizedClose); - - if (_allowDiskUse) { - tassert(7039549, - "disk use enabled for HashAggStage but incorrect number of merging expresssions", - _aggs.size() == _mergingExprs.size()); - } - - if (_forceIncreasedSpilling) { - tassert(7039554, "'forceIncreasedSpilling' set but disk use not allowed", _allowDiskUse); - } } std::unique_ptr<PlanStage> HashAggStage::clone() const { - SlotExprPairVector aggs; - aggs.reserve(_aggs.size()); + value::SlotMap<std::unique_ptr<EExpression>> aggs; for (auto& [k, v] : _aggs) { - aggs.push_back({k, v->clone()}); + aggs.emplace(k, v->clone()); } - - SlotExprPairVector mergingExprsClone; - mergingExprsClone.reserve(_mergingExprs.size()); - for (auto&& [k, v] : _mergingExprs) { - mergingExprsClone.push_back({k, v->clone()}); - } - return std::make_unique<HashAggStage>(_children[0]->clone(), _gbs, std::move(aggs), @@ -95,34 +75,24 @@ std::unique_ptr<PlanStage> HashAggStage::clone() const { _optimizedClose, _collatorSlot, _allowDiskUse, - std::move(mergingExprsClone), - _commonStats.nodeId, - _forceIncreasedSpilling); + _commonStats.nodeId); } void HashAggStage::doSaveState(bool relinquishCursor) { if (relinquishCursor) { if (_rsCursor) { - _recordStore->saveCursor(_opCtx, _rsCursor); + _rsCursor->save(); } } if (_rsCursor) { _rsCursor->setSaveStorageCursorOnDetachFromOperationContext(!relinquishCursor); } - - if (_recordStore) { - _recordStore->saveState(); - } } void HashAggStage::doRestoreState(bool relinquishCursor) { invariant(_opCtx); - if (_recordStore) { - _recordStore->restoreState(); - } - if (_rsCursor && relinquishCursor) { - auto couldRestore = _recordStore->restoreCursor(_opCtx, _rsCursor); + auto couldRestore = _rsCursor->restore(); uassert(6196500, "HashAggStage could not restore cursor", couldRestore); } } @@ -150,36 +120,34 @@ void HashAggStage::prepare(CompileCtx& ctx) { } value::SlotSet dupCheck; - auto throwIfDupSlot = [&dupCheck](value::SlotId slot) { - auto [_, inserted] = dupCheck.emplace(slot); - tassert(7039551, "duplicate slot id", inserted); - }; - size_t counter = 0; // Process group by columns. for (auto& slot : _gbs) { - throwIfDupSlot(slot); + auto [it, inserted] = dupCheck.emplace(slot); + uassert(4822827, str::stream() << "duplicate field: " << slot, inserted); _inKeyAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); - // Construct accessors for obtaining the key values from either the hash table '_ht' or the - // '_recordStore'. + // Construct accessors for the key to be processed from either the '_ht' or the + // '_recordStore'. Before the memory limit is reached the '_outHashKeyAccessors' will carry + // the group-by keys, otherwise the '_outRecordStoreKeyAccessors' will carry the group-by + // keys. _outHashKeyAccessors.emplace_back(std::make_unique<HashKeyAccessor>(_htIt, counter)); _outRecordStoreKeyAccessors.emplace_back( - std::make_unique<value::MaterializedSingleRowAccessor>(_outKeyRowRecordStore, counter)); + std::make_unique<value::MaterializedSingleRowAccessor>(_aggKeyRecordStore, counter)); + + counter++; - // A 'SwitchAccessor' is used to point the '_outKeyAccessors' to the key coming from the - // '_ht' or the '_recordStore' when draining the HashAgg stage in getNext(). If no spilling - // occurred, the keys will be obtained from the hash table. If spilling kicked in, then all - // of the data is written out to the record store, so the 'SwitchAccessor' is reconfigured - // to obtain all of the keys from the spill table. + // A SwitchAccessor is used to point the '_outKeyAccessors' to the key coming from the '_ht' + // or the '_recordStore' when draining the HashAgg stage in getNext. The group-by key will + // either be in the '_ht' or the '_recordStore' if the key lives in memory, or if the key + // has been spilled to disk, respectively. The SwitchAccessor allows toggling between the + // two so the parent stage can read it through the '_outAccessors'. _outKeyAccessors.emplace_back( std::make_unique<value::SwitchAccessor>(std::vector<value::SlotAccessor*>{ _outHashKeyAccessors.back().get(), _outRecordStoreKeyAccessors.back().get()})); _outAccessors[slot] = _outKeyAccessors.back().get(); - - ++counter; } // Process seek keys (if any). The keys must come from outside of the subtree (by definition) so @@ -190,18 +158,26 @@ void HashAggStage::prepare(CompileCtx& ctx) { counter = 0; for (auto& [slot, expr] : _aggs) { - throwIfDupSlot(slot); - - // Just like with the output accessors for the keys, we construct output accessors for the - // aggregate values that read from either the hash table '_ht' or the '_recordStore'. + auto [it, inserted] = dupCheck.emplace(slot); + // Some compilers do not allow to capture local bindings by lambda functions (the one + // is used implicitly in uassert below), so we need a local variable to construct an + // error message. + const auto slotId = slot; + uassert(4822828, str::stream() << "duplicate field: " << slotId, inserted); + + // Construct accessors for the agg state to be processed from either the '_ht' or the + // '_recordStore' by the SwitchAccessor owned by '_outAggAccessors' below. _outRecordStoreAggAccessors.emplace_back( - std::make_unique<value::MaterializedSingleRowAccessor>(_outAggRowRecordStore, counter)); + std::make_unique<value::MaterializedSingleRowAccessor>(_aggValueRecordStore, counter)); _outHashAggAccessors.emplace_back(std::make_unique<HashAggAccessor>(_htIt, counter)); - - // A 'SwitchAccessor' is used to toggle the '_outAggAccessors' between the '_ht' and the - // '_recordStore'. Just like the key values, the aggregate values are always obtained from - // the hash table if no spilling occurred and are always obtained from the record store if - // spilling occurred. + counter++; + + // A SwitchAccessor is used to toggle the '_outAggAccessors' between the '_ht' and the + // '_recordStore' when updating the agg state via the bytecode. By compiling the agg + // EExpressions with a SwitchAccessor we can load the agg value into the memory of + // '_aggValueRecordStore' if the value comes from the '_recordStore' or we can use the + // agg value referenced through '_htIt' and run the bytecode to mutate the value through the + // SwitchAccessor. _outAggAccessors.emplace_back( std::make_unique<value::SwitchAccessor>(std::vector<value::SlotAccessor*>{ _outHashAggAccessors.back().get(), _outRecordStoreAggAccessors.back().get()})); @@ -211,32 +187,10 @@ void HashAggStage::prepare(CompileCtx& ctx) { ctx.root = this; ctx.aggExpression = true; ctx.accumulator = _outAggAccessors.back().get(); + _aggCodes.emplace_back(expr->compile(ctx)); ctx.aggExpression = false; - - ++counter; - } - - // If disk use is allowed, then we need to compile the merging expressions as well. - if (_allowDiskUse) { - counter = 0; - for (auto&& [spillSlot, mergingExpr] : _mergingExprs) { - throwIfDupSlot(spillSlot); - - _spilledAggsAccessors.push_back( - std::make_unique<value::MaterializedSingleRowAccessor>(_spilledAggRow, counter)); - _spilledAggsAccessorMap[spillSlot] = _spilledAggsAccessors.back().get(); - - ctx.root = this; - ctx.aggExpression = true; - ctx.accumulator = _outAggAccessors[counter].get(); - _mergingExprCodes.emplace_back(mergingExpr->compile(ctx)); - ctx.aggExpression = false; - - ++counter; - } } - _compiled = true; } @@ -246,15 +200,6 @@ value::SlotAccessor* HashAggStage::getAccessor(CompileCtx& ctx, value::SlotId sl return it->second; } } else { - // The slots into which we read spilled partial aggregates, accessible via - // '_spilledAggsAccessors', should only be visible to this stage. They are used internally - // when merging spilled partial aggregates and should never be read by ancestor stages. - // Therefore, they are only made visible when this stage is in the process of compiling - // itself. - if (auto it = _spilledAggsAccessorMap.find(slot); it != _spilledAggsAccessorMap.end()) { - return it->second; - } - return _children[0]->getAccessor(ctx, slot); } @@ -270,117 +215,99 @@ void HashAggStage::makeTemporaryRecordStore() { "No storage engine so HashAggStage cannot spill to disk", _opCtx->getServiceContext()->getStorageEngine()); assertIgnorePrepareConflictsBehavior(_opCtx); - _recordStore = std::make_unique<SpillingStore>(_opCtx); + _recordStore = _opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore( + _opCtx, KeyFormat::String); _specificStats.usedDisk = true; } void HashAggStage::spillRowToDisk(const value::MaterializedRow& key, const value::MaterializedRow& val) { - CollatorInterface* collator = nullptr; - if (_collatorAccessor) { - auto [colTag, colVal] = _collatorAccessor->getViewOfValue(); - collator = value::getCollatorView(colVal); - } - KeyString::Builder kb{KeyString::Version::kLatestVersion}; - // Serialize the key that will be used as the record id (rid) when storing the record in the - // record store. Use a keystring for the spilled entry's rid such that partial aggregates are - // guaranteed to have identical keystrings when their keys are equal with respect to the - // collation. - key.serializeIntoKeyString(kb, collator); - // Add a unique integer to the end of the key, since record ids must be unique. We want equal - // keys to be adjacent in the 'RecordStore' so that we can merge the partial aggregates with a - // single pass. - kb.appendNumberLong(_ridCounter++); - auto rid = RecordId(kb.getBuffer(), kb.getSize()); + key.serializeIntoKeyString(kb); + auto typeBits = kb.getTypeBits(); - if (collator) { - // The keystring cannot always be deserialized back to the original keys when a collation is - // in use, so we also store the unmodified key in the data part of the spilled record. - _recordStore->upsertToRecordStore(_opCtx, rid, key, val, false /*update*/); - } else { - auto typeBits = kb.getTypeBits(); - _recordStore->upsertToRecordStore(_opCtx, rid, val, typeBits, false /*update*/); - } + auto rid = RecordId(kb.getBuffer(), kb.getSize()); + boost::optional<value::MaterializedRow> valFromRs = + readFromRecordStore(_opCtx, _recordStore->rs(), rid); + tassert(6031100, "Spilling a row doesn't support updating it in the store.", !valFromRs); - _specificStats.spilledRecords++; + spillValueToDisk(rid, val, typeBits, false /*update*/); } -void HashAggStage::spill(MemoryCheckData& mcd) { - uassert(ErrorCodes::QueryExceededMemoryLimitNoDiskUseAllowed, - "Exceeded memory limit for $group, but didn't allow external spilling;" - " pass allowDiskUse:true to opt in", - _allowDiskUse); - - // Since we flush the entire hash table to disk, we also clear any state related to estimating - // memory consumption. - mcd.reset(); - - if (!_recordStore) { - makeTemporaryRecordStore(); - } - - for (auto&& it : *_ht) { - spillRowToDisk(it.first, it.second); +void HashAggStage::spillValueToDisk(const RecordId& key, + const value::MaterializedRow& val, + const KeyString::TypeBits& typeBits, + bool update) { + auto nBytes = upsertToRecordStore(_opCtx, _recordStore->rs(), key, val, typeBits, update); + if (!update) { + _specificStats.spilledRecords++; } - - auto& metricsCollector = ResourceConsumption::MetricsCollector::get(_opCtx); - // We're not actually doing any sorting here or using the 'Sorter' class, but for the purposes - // of $operationMetrics we incorporate the number of spilled records into the "keysSorted" - // metric. Similarly, "sorterSpills" despite the name counts the number of individual spill - // events. - metricsCollector.incrementKeysSorted(_ht->size()); - metricsCollector.incrementSorterSpills(1); - - _ht->clear(); - - ++_specificStats.numSpills; + _specificStats.lastSpilledRecordSize = nBytes; } // Checks memory usage. Ideally, we'd want to know the exact size of already accumulated data, but // we cannot, so we estimate it based on the last updated/inserted row, if we have one, or the first // row in the '_ht' table. If the estimated memory usage exceeds the allowed, this method initiates -// spilling. +// spilling (if haven't been done yet) and evicts some records from the '_ht' table into the temp +// store to keep the memory usage under the limit. void HashAggStage::checkMemoryUsageAndSpillIfNecessary(MemoryCheckData& mcd) { - invariant(!_ht->empty()); + // The '_ht' table might become empty in the degenerate case when all rows had to be evicted to + // meet the memory constraint during a previous check -- we don't need to keep checking memory + // usage in this case because the table will never get new rows. + if (_ht->empty()) { + return; + } // If the group-by key is empty we will only ever aggregate into a single row so no sense in - // spilling. + // spilling since we will just be moving a single row back and forth from disk to main memory. if (_inKeyAccessors.size() == 0) { return; } mcd.memoryCheckpointCounter++; - if (mcd.memoryCheckpointCounter < mcd.nextMemoryCheckpoint) { - // We haven't reached the next checkpoint at which we estimate memory usage and decide if we - // should spill. - return; - } - - const long estimatedRowSize = - _htIt->first.memUsageForSorter() + _htIt->second.memUsageForSorter(); - long long estimatedTotalSize = _ht->size() * estimatedRowSize; - - if (estimatedTotalSize >= _approxMemoryUseInBytesBeforeSpill) { - spill(mcd); - } else { - // Calculate the next memory checkpoint. We estimate it based on the prior growth of the - // '_ht' and the remaining available memory. If 'estimatedGainPerChildAdvance' suggests that - // the hash table is growing, then the checkpoint is estimated as some configurable - // percentage of the number of additional input rows that we would have to process to - // consume the remaining memory. On the other hand, a value of 'estimtedGainPerChildAdvance' - // close to zero indicates a stable hash stable size, in which case we can delay the next - // check progressively. + if (mcd.memoryCheckpointCounter >= mcd.nextMemoryCheckpoint) { + if (_htIt == _ht->end()) { + _htIt = _ht->begin(); + } + const long estimatedRowSize = + _htIt->first.memUsageForSorter() + _htIt->second.memUsageForSorter(); + long long estimatedTotalSize = _ht->size() * estimatedRowSize; const double estimatedGainPerChildAdvance = (static_cast<double>(estimatedTotalSize - mcd.lastEstimatedMemoryUsage) / mcd.memoryCheckpointCounter); + if (estimatedTotalSize >= _approxMemoryUseInBytesBeforeSpill) { + uassert(ErrorCodes::QueryExceededMemoryLimitNoDiskUseAllowed, + "Exceeded memory limit for $group, but didn't allow external spilling." + " Pass allowDiskUse:true to opt in.", + _allowDiskUse); + if (!_recordStore) { + makeTemporaryRecordStore(); + } + + // Evict enough rows into the temporary store to drop below the memory constraint. + const long rowsToEvictCount = + 1 + (estimatedTotalSize - _approxMemoryUseInBytesBeforeSpill) / estimatedRowSize; + for (long i = 0; !_ht->empty() && i < rowsToEvictCount; i++) { + spillRowToDisk(_htIt->first, _htIt->second); + _ht->erase(_htIt); + _htIt = _ht->begin(); + } + estimatedTotalSize = _ht->size() * estimatedRowSize; + } + + // Calculate the next memory checkpoint. We estimate it based on the prior growth of the + // '_ht' and the remaining available memory. We have to keep doing this even after starting + // to spill because some accumulators can grow in size inside '_ht' (with no bounds). + // Value of 'estimatedGainPerChildAdvance' can be negative if the previous checkpoint + // evicted any records. And a value close to zero indicates a stable size of '_ht' so can + // delay the next check progressively. const long nextCheckpointCandidate = (estimatedGainPerChildAdvance > 0.1) ? mcd.checkpointMargin * (_approxMemoryUseInBytesBeforeSpill - estimatedTotalSize) / estimatedGainPerChildAdvance - : mcd.nextMemoryCheckpoint * 2; - + : (estimatedGainPerChildAdvance < -0.1) ? mcd.atMostCheckFrequency + : mcd.nextMemoryCheckpoint * 2; mcd.nextMemoryCheckpoint = std::min<long>(mcd.memoryCheckFrequency, std::max<long>(mcd.atMostCheckFrequency, nextCheckpointCandidate)); @@ -406,30 +333,17 @@ void HashAggStage::open(bool reOpen) { 5402503, "collatorSlot must be of collator type", tag == value::TypeTags::collator); auto collatorView = value::getCollatorView(collatorVal); const value::MaterializedRowHasher hasher(collatorView); - _keyEq = value::MaterializedRowEq(collatorView); - _ht.emplace(0, hasher, _keyEq); + const value::MaterializedRowEq equator(collatorView); + _ht.emplace(0, hasher, equator); } else { _ht.emplace(); } _seekKeys.resize(_seekKeysAccessors.size()); - // Reset state since this stage may have been previously opened. - for (auto&& accessor : _outKeyAccessors) { - accessor->setIndex(0); - } - for (auto&& accessor : _outAggAccessors) { - accessor->setIndex(0); - } - if (_recordStore) { - _recordStore->resetCursor(_opCtx, _rsCursor); - } - _recordStore.reset(); - _outKeyRowRecordStore = {0}; - _outAggRowRecordStore = {0}; - _spilledAggRow = {0}; - _stashedNextRow = {0, 0}; - + // A default value for spilling a key to the record store. + value::MaterializedRow defaultVal{_outAggAccessors.size()}; + bool updateAggStateHt = false; MemoryCheckData memoryCheckData; while (_children[0]->getNext() == PlanState::ADVANCED) { @@ -441,35 +355,57 @@ void HashAggStage::open(bool reOpen) { key.reset(idx++, false, tag, val); } - bool newKey = false; - _htIt = _ht->find(key); - if (_htIt == _ht->end()) { - // The key is not present in the hash table yet, so we insert it and initialize the - // corresponding accumulator. Note that as a future optimization, we could avoid - // doing a lookup both in the 'find()' call and in 'emplace()'. - newKey = true; + + if (_htIt = _ht->find(key); !_recordStore && _htIt == _ht->end()) { + // The memory limit hasn't been reached yet, insert a new key in '_ht' by copying + // the key. Note as a future optimization, we should avoid the lookup in the find() + // call and the emplace. key.makeOwned(); auto [it, _] = _ht->emplace(std::move(key), value::MaterializedRow{0}); + // Initialize accumulators. it->second.resize(_outAggAccessors.size()); _htIt = it; } - - // Accumulate state in '_ht'. - for (size_t idx = 0; idx < _outAggAccessors.size(); ++idx) { - auto [owned, tag, val] = _bytecode.run(_aggCodes[idx].get()); - _outHashAggAccessors[idx]->reset(owned, tag, val); - } - - if (_forceIncreasedSpilling && !newKey) { - // If configured to spill more than usual, we spill after seeing the same key twice. - spill(memoryCheckData); + updateAggStateHt = _htIt != _ht->end(); + + if (updateAggStateHt) { + // Accumulate state in '_ht' by pointing the '_outAggAccessors' the + // '_outHashAggAccessors'. + for (size_t idx = 0; idx < _outAggAccessors.size(); ++idx) { + _outAggAccessors[idx]->setIndex(0); + auto [owned, tag, val] = _bytecode.run(_aggCodes[idx].get()); + _outHashAggAccessors[idx]->reset(owned, tag, val); + } } else { - // Estimates how much memory is being used. If we estimate that the hash table - // exceeds the allotted memory budget, its contents are spilled to the - // '_recordStore' and '_ht' is cleared. - checkMemoryUsageAndSpillIfNecessary(memoryCheckData); + // The memory limit has been reached and the key wasn't in the '_ht' so we need + // to spill it to the '_recordStore'. + KeyString::Builder kb{KeyString::Version::kLatestVersion}; + // 'key' is moved only when 'updateAggStateHt' ends up "true", so it's safe to + // ignore the warning. + key.serializeIntoKeyString(kb); // NOLINT(bugprone-use-after-move) + auto typeBits = kb.getTypeBits(); + + auto rid = RecordId(kb.getBuffer(), kb.getSize()); + + boost::optional<value::MaterializedRow> valFromRs = + readFromRecordStore(_opCtx, _recordStore->rs(), rid); + if (!valFromRs) { + _aggValueRecordStore = defaultVal; + } else { + _aggValueRecordStore = *valFromRs; + } + + for (size_t idx = 0; idx < _outAggAccessors.size(); ++idx) { + _outAggAccessors[idx]->setIndex(1); + auto [owned, tag, val] = _bytecode.run(_aggCodes[idx].get()); + _aggValueRecordStore.reset(idx, owned, tag, val); + } + spillValueToDisk(rid, _aggValueRecordStore, typeBits, valFromRs ? true : false); } + // Estimates how much memory is being used and might start spilling. + checkMemoryUsageAndSpillIfNecessary(memoryCheckData); + if (_tracker && _tracker->trackProgress<TrialRunTracker::kNumResults>(1)) { // During trial runs, we want to limit the amount of work done by opening a blocking // stage, like this one. The blocking stage tracks the number of documents it has @@ -485,34 +421,9 @@ void HashAggStage::open(bool reOpen) { _children[0]->close(); _childOpened = false; } - - // If we spilled at any point while consuming the input, then do one final spill to write - // any leftover contents of '_ht' to the record store. That way, when recovering the input - // from the record store and merging partial aggregates we don't have to worry about the - // possibility of some of the data being in the hash table and some being in the record - // store. - if (_recordStore) { - if (!_ht->empty()) { - spill(memoryCheckData); - } - - _specificStats.spilledDataStorageSize = _recordStore->rs()->storageSize(_opCtx); - - // Establish a cursor, positioned at the beginning of the record store. - _rsCursor = _recordStore->getCursor(_opCtx); - - // Callers will be obtaining the results from the spill table, so set the - // 'SwitchAccessors' so that they refer to the rows recovered from the record store - // under the hood. - for (auto&& accessor : _outKeyAccessors) { - accessor->setIndex(1); - } - for (auto&& accessor : _outAggAccessors) { - accessor->setIndex(1); - } - } } + if (!_seekKeysAccessors.empty()) { // Copy keys in order to do the lookup. size_t idx = 0; @@ -523,103 +434,20 @@ void HashAggStage::open(bool reOpen) { } _htIt = _ht->end(); -} - -HashAggStage::SpilledRow HashAggStage::deserializeSpilledRecord(const Record& record, - BufBuilder& keyBuffer) { - // Read the values and type bits out of the value part of the record. - BufReader valReader(record.data.data(), record.data.size()); - - auto val = value::MaterializedRow::deserializeForSorter(valReader, {nullptr /*collator*/}); - auto typeBits = KeyString::TypeBits::fromBuffer(KeyString::Version::kLatestVersion, &valReader); - - keyBuffer.reset(); - auto key = value::MaterializedRow::deserializeFromKeyString( - decodeKeyString(record.id, typeBits), &keyBuffer, _gbs.size() /*numPrefixValuesToRead*/); - return {std::move(key), std::move(val)}; -} - -HashAggStage::SpilledRow HashAggStage::deserializeSpilledRecord(const Record& record, - const CollatorInterface& collator) { - BufReader valReader(record.data.data(), record.data.size()); - // When a collator has been defined, both the key and the value are stored in the data part of - // the record. First read the key and then read the value. - auto key = value::MaterializedRow::deserializeForSorter(valReader, {&collator}); - auto val = value::MaterializedRow::deserializeForSorter(valReader, {&collator}); - return {std::move(key), std::move(val)}; -} - -PlanState HashAggStage::getNextSpilled() { - CollatorInterface* collator = nullptr; - if (_collatorAccessor) { - auto [colTag, colVal] = _collatorAccessor->getViewOfValue(); - collator = value::getCollatorView(colVal); - } - - // Use the appropriate method to deserialize the record based on whether a collator is used. - auto recoverSpilledRecord = - [this](const Record& record, BufBuilder& keyBuffer, const CollatorInterface* collator) { - if (collator) { - return deserializeSpilledRecord(record, *collator); - } - return deserializeSpilledRecord(record, keyBuffer); - }; - - if (_stashedNextRow.first.isEmpty()) { - auto nextRecord = _rsCursor->next(); - if (!nextRecord) { - return trackPlanState(PlanState::IS_EOF); - } - - // We are just starting the process of merging the spilled file segments. - auto recoveredRow = recoverSpilledRecord(*nextRecord, _outKeyRowRSBuffer, collator); - - _outKeyRowRecordStore = std::move(recoveredRow.first); - _outAggRowRecordStore = std::move(recoveredRow.second); - } else { - // We peeked at the next key last time around. - _outKeyRowRSBuffer = std::move(_stashedKeyBuffer); - _outKeyRowRecordStore = std::move(_stashedNextRow.first); - _outAggRowRecordStore = std::move(_stashedNextRow.second); - // Clear the stashed row. - _stashedNextRow = {0, 0}; - } - - // Find additional partial aggregates for the same key and merge them in order to compute the - // final output. - for (auto nextRecord = _rsCursor->next(); nextRecord; nextRecord = _rsCursor->next()) { - auto recoveredRow = recoverSpilledRecord(*nextRecord, _stashedKeyBuffer, collator); - if (!_keyEq(recoveredRow.first, _outKeyRowRecordStore)) { - // The newly recovered spilled row belongs to a new key, so we're done merging partial - // aggregates for the old key. Save the new row for later and return advanced. - _stashedNextRow = std::move(recoveredRow); - return trackPlanState(PlanState::ADVANCED); - } - - // Merge in the new partial aggregate values. - _spilledAggRow = std::move(recoveredRow.second); - for (size_t idx = 0; idx < _mergingExprCodes.size(); ++idx) { - auto [owned, tag, val] = _bytecode.run(_mergingExprCodes[idx].get()); - _outRecordStoreAggAccessors[idx]->reset(owned, tag, val); - } + // Set the SwitchAccessors to point to the '_ht' so we can drain it first before draining the + // '_recordStore' in getNext(). + for (size_t idx = 0; idx < _outAggAccessors.size(); ++idx) { + _outAggAccessors[idx]->setIndex(0); } - - return trackPlanState(PlanState::ADVANCED); + _drainingRecordStore = false; } PlanState HashAggStage::getNext() { auto optTimer(getOptTimer(_opCtx)); - // If we've spilled, then we need to produce the output by merging the spilled segments from the - // spill file. - if (_recordStore) { - return getNextSpilled(); - } - - // We didn't spill. Obtain the next output row from the hash table. - if (_htIt == _ht->end()) { - // First invocation of getNext() after open(). + if (_htIt == _ht->end() && !_drainingRecordStore) { + // First invocation of getNext() after open() when not draining the '_recordStore'. if (!_seekKeysAccessors.empty()) { _htIt = _ht->find(_seekKeys); } else { @@ -628,13 +456,53 @@ PlanState HashAggStage::getNext() { } else if (!_seekKeysAccessors.empty()) { // Subsequent invocation with seek keys. Return only 1 single row (if any). _htIt = _ht->end(); - } else { + } else if (!_drainingRecordStore) { + // Returning the results of the entire hash table first before draining the '_recordStore'. ++_htIt; } - if (_htIt == _ht->end()) { - // The hash table has been drained (and we never spilled to disk) so we're done. + if (_htIt == _ht->end() && !_recordStore) { + // The hash table has been drained and nothing was spilled to disk. return trackPlanState(PlanState::IS_EOF); + } else if (_htIt != _ht->end()) { + // Drain the '_ht' on the next 'getNext()' call. + return trackPlanState(PlanState::ADVANCED); + } else if (_seekKeysAccessors.empty()) { + // A record store was created to spill to disk. Drain it then clean it up. + if (!_rsCursor) { + _rsCursor = _recordStore->rs()->getCursor(_opCtx); + } + auto nextRecord = _rsCursor->next(); + if (nextRecord) { + // Point the out accessors to the recordStore accessors to allow parent stages to read + // the agg state from the '_recordStore'. + if (!_drainingRecordStore) { + for (size_t i = 0; i < _outKeyAccessors.size(); ++i) { + _outKeyAccessors[i]->setIndex(1); + } + for (size_t i = 0; i < _outAggAccessors.size(); ++i) { + _outAggAccessors[i]->setIndex(1); + } + } + _drainingRecordStore = true; + + // Read the agg state value from the '_recordStore' and Reconstruct the key from the + // typeBits stored along side of the value. + BufReader valReader(nextRecord->data.data(), nextRecord->data.size()); + auto val = value::MaterializedRow::deserializeForSorter(valReader, {}); + auto typeBits = + KeyString::TypeBits::fromBuffer(KeyString::Version::kLatestVersion, &valReader); + _aggValueRecordStore = val; + + _aggKeyRSBuffer.reset(); + _aggKeyRecordStore = value::MaterializedRow::deserializeFromKeyString( + decodeKeyString(nextRecord->id, typeBits), &_aggKeyRSBuffer); + return trackPlanState(PlanState::ADVANCED); + } else { + _rsCursor.reset(); + _recordStore.reset(); + return trackPlanState(PlanState::IS_EOF); + } } else { return trackPlanState(PlanState::ADVANCED); } @@ -654,19 +522,11 @@ std::unique_ptr<PlanStageStats> HashAggStage::getStats(bool includeDebugInfo) co childrenBob.append(str::stream() << slot, printer.print(expr->debugPrint())); } } - - if (!_mergingExprs.empty()) { - BSONObjBuilder nestedBuilder{bob.subobjStart("mergingExprs")}; - for (auto&& [slot, expr] : _mergingExprs) { - nestedBuilder.append(str::stream() << slot, printer.print(expr->debugPrint())); - } - } - // Spilling stats. bob.appendBool("usedDisk", _specificStats.usedDisk); - bob.appendNumber("numSpills", _specificStats.numSpills); bob.appendNumber("spilledRecords", _specificStats.spilledRecords); - bob.appendNumber("spilledDataStorageSize", _specificStats.spilledDataStorageSize); + bob.appendNumber("spilledBytesApprox", + _specificStats.lastSpilledRecordSize * _specificStats.spilledRecords); ret->debugInfo = bob.obj(); } @@ -684,15 +544,11 @@ void HashAggStage::close() { trackClose(); _ht = boost::none; - if (_recordStore && _opCtx) { - _recordStore->resetCursor(_opCtx, _rsCursor); + if (_recordStore) { + // A record store was created to spill to disk. Clean it up. + _recordStore.reset(); + _drainingRecordStore = false; } - _rsCursor.reset(); - _recordStore.reset(); - _outKeyRowRecordStore = {0}; - _outAggRowRecordStore = {0}; - _spilledAggRow = {0}; - _stashedNextRow = {0, 0}; if (_childOpened) { _children[0]->close(); @@ -715,7 +571,7 @@ std::vector<DebugPrinter::Block> HashAggStage::debugPrint() const { ret.emplace_back(DebugPrinter::Block("[`")); bool first = true; - for (auto&& [slot, expr] : _aggs) { + value::orderedSlotMapTraverse(_aggs, [&](auto slot, auto&& expr) { if (!first) { ret.emplace_back(DebugPrinter::Block("`,")); } @@ -724,7 +580,7 @@ std::vector<DebugPrinter::Block> HashAggStage::debugPrint() const { ret.emplace_back("="); DebugPrinter::addBlocks(ret, expr->debugPrint()); first = false; - } + }); ret.emplace_back("`]"); if (!_seekKeysSlots.empty()) { @@ -739,28 +595,6 @@ std::vector<DebugPrinter::Block> HashAggStage::debugPrint() const { ret.emplace_back("`]"); } - if (!_mergingExprs.empty()) { - ret.emplace_back("spillSlots[`"); - for (size_t idx = 0; idx < _mergingExprs.size(); ++idx) { - if (idx) { - ret.emplace_back("`,"); - } - - DebugPrinter::addIdentifier(ret, _mergingExprs[idx].first); - } - ret.emplace_back("`]"); - - ret.emplace_back("mergingExprs[`"); - for (size_t idx = 0; idx < _mergingExprs.size(); ++idx) { - if (idx) { - ret.emplace_back("`,"); - } - - DebugPrinter::addBlocks(ret, _mergingExprs[idx].second->debugPrint()); - } - ret.emplace_back("`]"); - } - if (!_optimizedClose) { ret.emplace_back("reopen"); } @@ -781,7 +615,6 @@ size_t HashAggStage::estimateCompileTimeSize() const { size += size_estimator::estimate(_gbs); size += size_estimator::estimate(_aggs); size += size_estimator::estimate(_seekKeysSlots); - size += size_estimator::estimate(_mergingExprs); return size; } diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.h b/src/mongo/db/exec/sbe/stages/hash_agg.h index 91f91051363..19fbca9d1c7 100644 --- a/src/mongo/db/exec/sbe/stages/hash_agg.h +++ b/src/mongo/db/exec/sbe/stages/hash_agg.h @@ -29,9 +29,10 @@ #pragma once +#include <unordered_map> + #include "mongo/db/exec/sbe/expressions/expression.h" #include "mongo/db/exec/sbe/stages/stages.h" -#include "mongo/db/exec/sbe/util/spilling.h" #include "mongo/db/exec/sbe/vm/vm.h" #include "mongo/db/query/query_knobs_gen.h" #include "mongo/db/storage/temporary_record_store.h" @@ -60,38 +61,21 @@ namespace sbe { * determining whether two group-by keys are equal. For instance, the plan may require us to do a * case-insensitive group on a string field. * - * The 'allowDiskUse' flag controls whether this stage can spill. If false and the memory budget is - * exhausted, this stage throws a query-fatal error with code - * 'QueryExceededMemoryLimitNoDiskUseAllowed'. If true, then spilling is possible and the caller - * must provide a vector of 'mergingExprs'. This is a vector of (slot, expression) pairs which is - * symmetrical with 'aggs'. The slots are only visible internally and are used to store partial - * aggregate values that have been recovered from the spill table. Each of the expressions is an agg - * function which merges the partial aggregate value from this slot into the final aggregate value. - * In the debug string output, the internal slots used to house the partial aggregates are printed - * as a list of "spillSlots" and the expressions are printed as a parallel list of "mergingExprs". - * - * If 'forcedIncreasedSpilling' is true, then this stage will spill frequently even if the memory - * limit is not reached. This is intended to be used in test contexts to exercise the otherwise - * infrequently used spilling logic. - * * Debug string representation: * - * group [<group by slots>] [slot_1 = expr_1, ..., slot_n = expr_n] [<seek slots>]? - * spillSlots[slot_1, ..., slot_n] mergingExprs[expr_1, ..., expr_n] reopen? collatorSlot? - * childStage + * group [<group by slots>] [slot_1 = expr_1, ..., slot_n = expr_n] [<seek slots>]? reopen? + * collatorSlot? childStage */ class HashAggStage final : public PlanStage { public: HashAggStage(std::unique_ptr<PlanStage> input, value::SlotVector gbs, - SlotExprPairVector aggs, + value::SlotMap<std::unique_ptr<EExpression>> aggs, value::SlotVector seekKeysSlots, bool optimizedClose, boost::optional<value::SlotId> collatorSlot, bool allowDiskUse, - SlotExprPairVector mergingExprs, - PlanNodeId planNodeId, - bool forceIncreasedSpilling = false); + PlanNodeId planNodeId); std::unique_ptr<PlanStage> clone() const final; @@ -124,147 +108,99 @@ private: using HashKeyAccessor = value::MaterializedRowKeyAccessor<TableType::iterator>; using HashAggAccessor = value::MaterializedRowValueAccessor<TableType::iterator>; - using SpilledRow = std::pair<value::MaterializedRow, value::MaterializedRow>; + void makeTemporaryRecordStore(); + + /** + * Spills a key and value pair to the '_recordStore' where the semantics are insert or update + * depending on the 'update' flag. When the 'update' flag is true this method already expects + * the 'key' to be inserted into the '_recordStore', otherwise the 'key' and 'val' pair are + * fresh. + * + * This method expects the key to be seralized into a KeyString::Value so that the key is + * memcmp-able and lookups can be done to update the 'val' in the '_recordStore'. Note that the + * 'typeBits' are needed to reconstruct the spilled 'key' when calling 'getNext' to deserialize + * the 'key' to a MaterializedRow. Since the '_recordStore' only stores the memcmp-able part of + * the KeyString we need to carry the 'typeBits' separately, and we do this by appending the + * 'typeBits' to the end of the serialized 'val' buffer and store them at the leaves of the + * backing B-tree of the '_recordStore'. used as the RecordId. + */ + void spillValueToDisk(const RecordId& key, + const value::MaterializedRow& val, + const KeyString::TypeBits& typeBits, + bool update); + void spillRowToDisk(const value::MaterializedRow& key, + const value::MaterializedRow& defaultVal); /** * We check amount of used memory every T processed incoming records, where T is calculated * based on the estimated used memory and its recent growth. When the memory limit is exceeded, - * 'checkMemoryUsageAndSpillIfNecessary()' will create '_recordStore' (if it hasn't already been - * created) and spill the contents of the hash table into this record store. + * 'checkMemoryUsageAndSpillIfNecessary()' will create '_recordStore' and might spill some of + * the already accumulated data into it. */ struct MemoryCheckData { - MemoryCheckData() { - reset(); - } - - void reset() { - memoryCheckFrequency = std::min(atMostCheckFrequency, atLeastMemoryCheckFrequency); - nextMemoryCheckpoint = 0; - memoryCheckpointCounter = 0; - lastEstimatedMemoryUsage = 0; - } - const double checkpointMargin = internalQuerySBEAggMemoryUseCheckMargin.load(); - const int64_t atMostCheckFrequency = internalQuerySBEAggMemoryCheckPerAdvanceAtMost.load(); - const int64_t atLeastMemoryCheckFrequency = + const long atMostCheckFrequency = internalQuerySBEAggMemoryCheckPerAdvanceAtMost.load(); + const long atLeastMemoryCheckFrequency = internalQuerySBEAggMemoryCheckPerAdvanceAtLeast.load(); // The check frequency upper bound, which start at 'atMost' and exponentially backs off // to 'atLeast' as more data is accumulated. If 'atLeast' is less than 'atMost', the memory // checks will be done every 'atLeast' incoming records. - int64_t memoryCheckFrequency = 1; + long memoryCheckFrequency = 1; // The number of incoming records to process before the next memory checkpoint. - int64_t nextMemoryCheckpoint = 0; + long nextMemoryCheckpoint = 0; // The counter of the incoming records between memory checkpoints. - int64_t memoryCheckpointCounter = 0; - - int64_t lastEstimatedMemoryUsage = 0; - }; + long memoryCheckpointCounter = 0; - /** - * Inserts a key and value pair to the '_recordStore'. They key is serialized to a - * 'KeyString::Value' which becomes the 'RecordId'. This makes the keys memcmp-able and ensures - * that the record store ends up sorted by the group-by keys. - * - * Note that the 'typeBits' are needed to reconstruct the spilled 'key' to a 'MaterializedRow', - * but are not necessary for comparison purposes. Therefore, we carry the type bits separately - * from the record id, instead appending them to the end of the serialized 'val' buffer. - */ - void spillRowToDisk(const value::MaterializedRow& key, const value::MaterializedRow& val); + long long lastEstimatedMemoryUsage = 0; + MemoryCheckData() { + memoryCheckFrequency = std::min(atMostCheckFrequency, atLeastMemoryCheckFrequency); + } + }; void checkMemoryUsageAndSpillIfNecessary(MemoryCheckData& mcd); - void spill(MemoryCheckData& mcd); - - /** - * Given a 'record' from the record store, decodes it into a pair of materialized rows (one for - * the group-by keys and another for the agg values). - * - * The given 'keyBuffer' is cleared, and then used to hold data (e.g. long strings and other - * values that can't be inlined) obtained by decoding the 'RecordId' keystring to a - * 'MaterializedRow'. The values in the resulting 'MaterializedRow' may be pointers into - * 'keyBuffer', so it is important that 'keyBuffer' outlive the row. - * - * This method is used when there is no collator. - */ - SpilledRow deserializeSpilledRecord(const Record& record, BufBuilder& keyBuffer); - - /** - * Given a 'record' from the record store and a 'collator', decodes it into a pair of - * materialized rows (one for the group-by key and another one for the agg value). - * Both the group-by key and the agg value are read from the data part of the record. - */ - SpilledRow deserializeSpilledRecord(const Record& record, const CollatorInterface& collator); - - PlanState getNextSpilled(); - - void makeTemporaryRecordStore(); const value::SlotVector _gbs; - const SlotExprPairVector _aggs; + const value::SlotMap<std::unique_ptr<EExpression>> _aggs; const boost::optional<value::SlotId> _collatorSlot; const bool _allowDiskUse; const value::SlotVector _seekKeysSlots; // When this operator does not expect to be reopened (almost always) then it can close the child // early. const bool _optimizedClose{true}; - - // Expressions used to merge partial aggregates that have been spilled to disk and their - // corresponding input slots. For example, imagine that this list contains a pair (s12, - // sum(s12)). This means that the partial aggregate values will be read into slot s12 after - // being recovered from the spill table and can be merged using the 'sum()' agg function. - // - // When disk use is allowed, this vector must have the same length as '_aggs'. - const SlotExprPairVector _mergingExprs; - - // When true, we spill frequently without reaching the memory limit. This allows us to exercise - // the spilling logic more often in test contexts. - const bool _forceIncreasedSpilling; - value::SlotAccessorMap _outAccessors; - - // Accessors used to obtain the values of the group by slots when reading the input from the - // child. std::vector<value::SlotAccessor*> _inKeyAccessors; - // This buffer stores values for '_outKeyRowRecordStore'; values in the '_outKeyRowRecordStore' - // can be pointers that point to data in this buffer. - BufBuilder _outKeyRowRSBuffer; - // Accessors for the key slots provided as output by this stage. The keys can either come from - // the hash table or recovered from a temporary record store. We use a 'SwitchAccessor' to - // switch between these two cases. + // Accessors for the key stored in '_ht', a SwitchAccessor is used so we can produce the key + // from either the '_ht' or the '_recordStore'. std::vector<std::unique_ptr<HashKeyAccessor>> _outHashKeyAccessors; - // Row of key values to output used when recovering spilled data from the record store. - value::MaterializedRow _outKeyRowRecordStore{0}; - std::vector<std::unique_ptr<value::MaterializedSingleRowAccessor>> _outRecordStoreKeyAccessors; std::vector<std::unique_ptr<value::SwitchAccessor>> _outKeyAccessors; - // Accessors for the output aggregate results. The aggregates can either come from the hash - // table or can be computed after merging partial aggregates spilled to a record store. We use a - // 'SwitchAccessor' to switch between these two cases. - std::vector<std::unique_ptr<HashAggAccessor>> _outHashAggAccessors; - // Row of agg values to output used when recovering spilled data from the record store. - value::MaterializedRow _outAggRowRecordStore{0}; + // Accessor for the agg state value stored in the '_recordStore' when data is spilled to disk. + value::MaterializedRow _aggKeyRecordStore{0}; + value::MaterializedRow _aggValueRecordStore{0}; + std::vector<std::unique_ptr<value::MaterializedSingleRowAccessor>> _outRecordStoreKeyAccessors; std::vector<std::unique_ptr<value::MaterializedSingleRowAccessor>> _outRecordStoreAggAccessors; - std::vector<std::unique_ptr<value::SwitchAccessor>> _outAggAccessors; + + // This buffer stores values for the spilled '_aggKeyRecordStore' that's loaded into memory from + // the '_recordStore'. Values in the '_aggKeyRecordStore' row are pointers that point to data in + // this buffer. + BufBuilder _aggKeyRSBuffer; std::vector<value::SlotAccessor*> _seekKeysAccessors; value::MaterializedRow _seekKeys; - // Bytecode which gets executed to aggregate incoming rows into the hash table. + // Accesors for the agg state in '_ht', a SwitchAccessor is used so we can produce the agg state + // from either the '_ht' or the '_recordStore' when draining the HashAgg stage. + std::vector<std::unique_ptr<value::SwitchAccessor>> _outAggAccessors; + std::vector<std::unique_ptr<HashAggAccessor>> _outHashAggAccessors; std::vector<std::unique_ptr<vm::CodeFragment>> _aggCodes; - // Bytecode for the merging expressions, executed if partial aggregates are spilled to a record - // store and need to be subsequently combined. - std::vector<std::unique_ptr<vm::CodeFragment>> _mergingExprCodes; // Only set if collator slot provided on construction. value::SlotAccessor* _collatorAccessor = nullptr; - // Function object which can be used to check whether two materialized rows of key values are - // equal. This comparison is collation-aware if the query has a non-simple collation. - value::MaterializedRowEq _keyEq; - boost::optional<TableType> _ht; TableType::iterator _htIt; @@ -276,31 +212,10 @@ private: // Memory tracking and spilling to disk. const long long _approxMemoryUseInBytesBeforeSpill = internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.load(); - - // A record store which is instantiated and written to in the case of spilling. - std::unique_ptr<SpillingStore> _recordStore; + std::unique_ptr<TemporaryRecordStore> _recordStore; + bool _drainingRecordStore{false}; std::unique_ptr<SeekableRecordCursor> _rsCursor; - // A monotically increasing counter used to ensure uniqueness of 'RecordId' values. When - // spilling, the key is encoding into the 'RecordId' of the '_recordStore'. Record ids must be - // unique by definition, but we might end up spilling multiple partial aggregates for the same - // key. We ensure uniqueness by appending a unique integer to the end of this key, which is - // simply ignored during deserialization. - int64_t _ridCounter = 0; - - // Partial aggregates that have been spilled are read into '_spilledAggRow' and read using - // '_spilledAggsAccessors' so that they can be merged to compute the final aggregate value. - value::MaterializedRow _spilledAggRow{0}; - std::vector<std::unique_ptr<value::MaterializedSingleRowAccessor>> _spilledAggsAccessors; - value::SlotAccessorMap _spilledAggsAccessorMap; - - // Buffer to hold data for the deserialized key values from '_stashedNextRow'. - BufBuilder _stashedKeyBuffer; - // Place to stash the next keys and values during the streaming phase. The record store cursor - // doesn't offer a "peek" API, so we need to hold onto the next row between getNext() calls when - // the key value advances. - SpilledRow _stashedNextRow; - HashAggStats _specificStats; // If provided, used during a trial run to accumulate certain execution stats. Once the trial diff --git a/src/mongo/db/exec/sbe/stages/hash_lookup.cpp b/src/mongo/db/exec/sbe/stages/hash_lookup.cpp index 02e95307c4b..0d976f1a875 100644 --- a/src/mongo/db/exec/sbe/stages/hash_lookup.cpp +++ b/src/mongo/db/exec/sbe/stages/hash_lookup.cpp @@ -187,22 +187,6 @@ value::SlotAccessor* HashLookupStage::getAccessor(CompileCtx& ctx, value::SlotId return outerChild()->getAccessor(ctx, slot); } } -void HashLookupStage::doSaveState(bool relinquishCursor) { - if (_recordStoreHt) { - _recordStoreHt->saveState(); - } - if (_recordStoreBuf) { - _recordStoreBuf->saveState(); - } -} -void HashLookupStage::doRestoreState(bool relinquishCursor) { - if (_recordStoreHt) { - _recordStoreHt->restoreState(); - } - if (_recordStoreBuf) { - _recordStoreBuf->restoreState(); - } -} void HashLookupStage::reset() { _ht = boost::none; @@ -275,7 +259,7 @@ void HashLookupStage::addHashTableEntry(value::SlotAccessor* keyAccessor, size_t auto val = std::vector<size_t>{valueIndex}; auto [tagKey, valKey] = keyAccessor->getViewOfValue(); - spillIndicesToRecordStore(_recordStoreHt.get(), tagKey, valKey, val); + spillIndicesToRecordStore(_recordStoreHt->rs(), tagKey, valKey, val); } } else { // The key is already present in '_ht' so the memory will only grow by one size_t. If we @@ -297,7 +281,7 @@ void HashLookupStage::addHashTableEntry(value::SlotAccessor* keyAccessor, size_t // Evict the hash table value. _computedTotalMemUsage -= htIt->second.size() * sizeof(size_t); htIt->second.push_back(valueIndex); - spillIndicesToRecordStore(_recordStoreHt.get(), tagKeyView, valKeyView, htIt->second); + spillIndicesToRecordStore(_recordStoreHt->rs(), tagKeyView, valKeyView, htIt->second); _ht->erase(htIt); } } @@ -313,15 +297,17 @@ void HashLookupStage::makeTemporaryRecordStore() { _opCtx->getServiceContext()->getStorageEngine()); assertIgnorePrepareConflictsBehavior(_opCtx); - _recordStoreBuf = std::make_unique<SpillingStore>(_opCtx, KeyFormat::Long); + _recordStoreBuf = _opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore( + _opCtx, KeyFormat::Long); - _recordStoreHt = std::make_unique<SpillingStore>(_opCtx, KeyFormat::String); + _recordStoreHt = _opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore( + _opCtx, KeyFormat::String); _specificStats.usedDisk = true; } void HashLookupStage::spillBufferedValueToDisk(OperationContext* opCtx, - SpillingStore* rs, + RecordStore* rs, size_t bufferIdx, const value::MaterializedRow& val) { auto rid = getValueRecordId(bufferIdx); @@ -329,7 +315,15 @@ void HashLookupStage::spillBufferedValueToDisk(OperationContext* opCtx, BufBuilder buf; val.serializeForSorter(buf); - rs->upsertToRecordStore(opCtx, rid, buf, false); + assertIgnorePrepareConflictsBehavior(opCtx); + WriteUnitOfWork wuow(opCtx); + + auto status = rs->insertRecord(opCtx, rid, buf.buf(), buf.len(), Timestamp{}); + wuow.commit(); + + tassert(6373906, + str::stream() << "Failed to write to disk because " << status.getStatus().reason(), + status.isOK()); _specificStats.spilledBuffRecords++; // Add size of record ID + size of buffer. @@ -340,14 +334,14 @@ void HashLookupStage::spillBufferedValueToDisk(OperationContext* opCtx, size_t HashLookupStage::bufferValueOrSpill(value::MaterializedRow& value) { size_t bufferIndex = _valueId; const long long newMemUsage = _computedTotalMemUsage + size_estimator::estimate(value); - if (!hasSpilledBufToDisk() && newMemUsage <= _memoryUseInBytesBeforeSpill) { + if (newMemUsage <= _memoryUseInBytesBeforeSpill) { _buffer.emplace_back(std::move(value)); _computedTotalMemUsage = newMemUsage; } else { if (!hasSpilledBufToDisk()) { makeTemporaryRecordStore(); } - spillBufferedValueToDisk(_opCtx, _recordStoreBuf.get(), bufferIndex, value); + spillBufferedValueToDisk(_opCtx, _recordStoreBuf->rs(), bufferIndex, value); } _valueId++; return bufferIndex; @@ -433,7 +427,7 @@ void HashLookupStage::accumulateFromValueIndices(const C& bufferIndices) { // We must shift the '_bufferIt' index by one when using it as a RecordId because a // RecordId of 0 is invalid. auto rid = getValueRecordId(_bufferIt); - auto rsValue = _recordStoreBuf->readFromRecordStore(_opCtx, rid); + auto rsValue = readFromRecordStore(_opCtx, _recordStoreBuf->rs(), rid); if (!rsValue) { tasserted(6373900, "bufferIdx not found in record store"); } @@ -449,7 +443,7 @@ void HashLookupStage::accumulateFromValueIndices(const C& bufferIndices) { } } -void HashLookupStage::writeIndicesToRecordStore(SpillingStore* rs, +void HashLookupStage::writeIndicesToRecordStore(RecordStore* rs, value::TypeTags tagKey, value::Value valKey, const std::vector<size_t>& value, @@ -464,7 +458,7 @@ void HashLookupStage::writeIndicesToRecordStore(SpillingStore* rs, key.reset(0, false, tagKey, valKey); auto [rid, typeBits] = serializeKeyForRecordStore(key); - rs->upsertToRecordStore(_opCtx, rid, buf, typeBits, update); + upsertToRecordStore(_opCtx, rs, rid, buf, typeBits, update); if (!update) { _specificStats.spilledHtRecords++; // Add the size of key (which comprises of the memory usage for the key + its type bits), @@ -477,7 +471,7 @@ void HashLookupStage::writeIndicesToRecordStore(SpillingStore* rs, } boost::optional<std::vector<size_t>> HashLookupStage::readIndicesFromRecordStore( - SpillingStore* rs, value::TypeTags tagKey, value::Value valKey) { + RecordStore* rs, value::TypeTags tagKey, value::Value valKey) { _probeKey.reset(0, false, tagKey, valKey); auto [rid, _] = serializeKeyForRecordStore(_probeKey); @@ -496,7 +490,7 @@ boost::optional<std::vector<size_t>> HashLookupStage::readIndicesFromRecordStore return boost::none; } -void HashLookupStage::spillIndicesToRecordStore(SpillingStore* rs, +void HashLookupStage::spillIndicesToRecordStore(RecordStore* rs, value::TypeTags tagKey, value::Value valKey, const std::vector<size_t>& value) { @@ -551,7 +545,7 @@ PlanState HashLookupStage::getNext() { normalizeStringIfCollator(tagElemView, valElemView); auto indicesFromRS = readIndicesFromRecordStore( - _recordStoreHt.get(), tagElemCollView, valElemCollView); + _recordStoreHt->rs(), tagElemCollView, valElemCollView); if (indicesFromRS) { indices.insert(indicesFromRS->begin(), indicesFromRS->end()); } @@ -573,7 +567,7 @@ PlanState HashLookupStage::getNext() { normalizeStringIfCollator(tagKeyView, valKeyView); auto indicesFromRS = readIndicesFromRecordStore( - _recordStoreHt.get(), tagKeyCollView, valKeyCollView); + _recordStoreHt->rs(), tagKeyCollView, valKeyCollView); if (indicesFromRS) { accumulateFromValueIndices(*indicesFromRS); } diff --git a/src/mongo/db/exec/sbe/stages/hash_lookup.h b/src/mongo/db/exec/sbe/stages/hash_lookup.h index b312e0a68f4..2e3f0b34816 100644 --- a/src/mongo/db/exec/sbe/stages/hash_lookup.h +++ b/src/mongo/db/exec/sbe/stages/hash_lookup.h @@ -33,7 +33,6 @@ #include "mongo/db/exec/sbe/expressions/expression.h" #include "mongo/db/exec/sbe/stages/stages.h" -#include "mongo/db/exec/sbe/util/spilling.h" #include "mongo/db/exec/sbe/vm/vm.h" #include "mongo/db/query/query_knobs_gen.h" @@ -102,10 +101,6 @@ public: std::vector<DebugPrinter::Block> debugPrint() const final; size_t estimateCompileTimeSize() const final; -protected: - void doSaveState(bool relinquishCursor) override; - void doRestoreState(bool relinquishCursor) override; - private: using HashTableType = std::unordered_map<value::MaterializedRow, // NOLINT std::vector<size_t>, @@ -124,23 +119,23 @@ private: // Spilling helpers. void addHashTableEntry(value::SlotAccessor* keyAccessor, size_t valueIndex); void spillBufferedValueToDisk(OperationContext* opCtx, - SpillingStore* rs, + RecordStore* rs, size_t bufferIdx, const value::MaterializedRow&); size_t bufferValueOrSpill(value::MaterializedRow& value); void setInnerProjectSwitchAccessor(int idx); - boost::optional<std::vector<size_t>> readIndicesFromRecordStore(SpillingStore* rs, + boost::optional<std::vector<size_t>> readIndicesFromRecordStore(RecordStore* rs, value::TypeTags tagKey, value::Value valKey); - void writeIndicesToRecordStore(SpillingStore* rs, + void writeIndicesToRecordStore(RecordStore* rs, value::TypeTags tagKey, value::Value valKey, const std::vector<size_t>& value, bool update); - void spillIndicesToRecordStore(SpillingStore* rs, + void spillIndicesToRecordStore(RecordStore* rs, value::TypeTags tagKey, value::Value valKey, const std::vector<size_t>& value); @@ -234,8 +229,8 @@ private: // rows in '_buffer'. long long _computedTotalMemUsage = 0; - std::unique_ptr<SpillingStore> _recordStoreHt; - std::unique_ptr<SpillingStore> _recordStoreBuf; + std::unique_ptr<TemporaryRecordStore> _recordStoreHt; + std::unique_ptr<TemporaryRecordStore> _recordStoreBuf; HashLookupStats _specificStats; }; diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp index fe88d9c9095..bfad6d9a2ae 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp +++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp @@ -45,7 +45,6 @@ IndexScanStage::IndexScanStage(UUID collUuid, boost::optional<value::SlotId> recordSlot, boost::optional<value::SlotId> recordIdSlot, boost::optional<value::SlotId> snapshotIdSlot, - boost::optional<value::SlotId> indexIdSlot, IndexKeysInclusionSet indexKeysToInclude, value::SlotVector vars, boost::optional<value::SlotId> seekKeySlotLow, @@ -59,7 +58,6 @@ IndexScanStage::IndexScanStage(UUID collUuid, _recordSlot(recordSlot), _recordIdSlot(recordIdSlot), _snapshotIdSlot(snapshotIdSlot), - _indexIdSlot(indexIdSlot), _indexKeysToInclude(indexKeysToInclude), _vars(std::move(vars)), _seekKeySlotLow(seekKeySlotLow), @@ -78,7 +76,6 @@ std::unique_ptr<PlanStage> IndexScanStage::clone() const { _recordSlot, _recordIdSlot, _snapshotIdSlot, - _indexIdSlot, _indexKeysToInclude, _vars, _seekKeySlotLow, @@ -131,17 +128,10 @@ void IndexScanStage::prepare(CompileCtx& ctx) { static_cast<bool>(entry)); _ordering = entry->ordering(); - auto [indexIdTag, indexIdVal] = value::makeNewString(StringData(_indexName)); - _indexIdAccessor.reset(indexIdTag, indexIdVal); - - if (_indexIdSlot) { - _indexIdViewAccessor.reset(indexIdTag, indexIdVal); - } else { - _indexIdViewAccessor.reset(); - } - if (_snapshotIdAccessor) { - _latestSnapshotId = _opCtx->recoveryUnit()->getSnapshotId().toNumber(); + _snapshotIdAccessor->reset( + value::TypeTags::NumberInt64, + value::bitcastFrom<uint64_t>(_opCtx->recoveryUnit()->getSnapshotId().toNumber())); } } @@ -158,10 +148,6 @@ value::SlotAccessor* IndexScanStage::getAccessor(CompileCtx& ctx, value::SlotId return _snapshotIdAccessor.get(); } - if (_indexIdSlot && *_indexIdSlot == slot) { - return &_indexIdViewAccessor; - } - if (auto it = _accessorMap.find(slot); it != _accessorMap.end()) { return it->second; } @@ -233,7 +219,9 @@ void IndexScanStage::doRestoreState(bool relinquishCursor) { // Yield is the only time during plan execution that the snapshotId can change. As such, we // update it accordingly as part of yield recovery. if (_snapshotIdAccessor) { - _latestSnapshotId = _opCtx->recoveryUnit()->getSnapshotId().toNumber(); + _snapshotIdAccessor->reset( + value::TypeTags::NumberInt64, + value::bitcastFrom<uint64_t>(_opCtx->recoveryUnit()->getSnapshotId().toNumber())); } } @@ -397,12 +385,6 @@ PlanState IndexScanStage::getNext() { false, value::TypeTags::RecordId, value::bitcastFrom<RecordId*>(&_nextRecord->loc)); } - if (_snapshotIdAccessor) { - // Copy the latest snapshot ID into the 'snapshotId' slot. - _snapshotIdAccessor->reset(value::TypeTags::NumberInt64, - value::bitcastFrom<uint64_t>(_latestSnapshotId)); - } - if (_accessors.size()) { _valuesBuffer.reset(); readKeyStringValueIntoAccessors( @@ -441,9 +423,6 @@ std::unique_ptr<PlanStageStats> IndexScanStage::getStats(bool includeDebugInfo) if (_snapshotIdSlot) { bob.appendNumber("snapshotIdSlot", static_cast<long long>(*_snapshotIdSlot)); } - if (_indexIdSlot) { - bob.appendNumber("indexIdSlot", static_cast<long long>(*_indexIdSlot)); - } if (_seekKeySlotLow) { bob.appendNumber("seekKeySlotLow", static_cast<long long>(*_seekKeySlotLow)); } @@ -492,12 +471,6 @@ std::vector<DebugPrinter::Block> IndexScanStage::debugPrint() const { DebugPrinter::addIdentifier(ret, DebugPrinter::kNoneKeyword); } - if (_indexIdSlot) { - DebugPrinter::addIdentifier(ret, _indexIdSlot.value()); - } else { - DebugPrinter::addIdentifier(ret, DebugPrinter::kNoneKeyword); - } - ret.emplace_back(DebugPrinter::Block("[`")); size_t varIndex = 0; for (size_t keyIndex = 0; keyIndex < _indexKeysToInclude.size(); ++keyIndex) { diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h index 33fcd352d0c..ce00ef17128 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.h +++ b/src/mongo/db/exec/sbe/stages/ix_scan.h @@ -51,8 +51,7 @@ namespace mongo::sbe { * The "output" slots are * - 'recordSlot': the "KeyString" representing the index entry, * - 'recordIdSlot': a reference that can be used to fetch the entire document, - * - 'snapshotIdSlot': the storage snapshot that this index scan is reading from, - * - 'indexIdSlot': the name of the index being read from, and + * - 'snapshotIdSlot': the storage snapshot that this index scan is reading from, and * - 'vars': one slot for each value in the index key that should be "projected" out of the entry. * * The 'indexKeysToInclude' bitset determines which values are included in the projection based @@ -64,11 +63,10 @@ namespace mongo::sbe { * * Debug string representation: * - * ixscan recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? - * [slot_1 = fieldNo_1, ..., slot2 = fieldNo_n] + * ixscan recordSlot? recordIdSlot? snapshotIdSlot? [slot_1 = fieldNo_1, ..., slot2 = fieldNo_n] * collectionUuid indexName forward * - * ixseek lowKey highKey recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? + * ixseek lowKey highKey recordSlot? recordIdSlot? snapshotIdSlot? * [slot_1 = fieldNo_1, ..., slot2 = fieldNo_n] * collectionUuid indexName forward */ @@ -80,7 +78,6 @@ public: boost::optional<value::SlotId> recordSlot, boost::optional<value::SlotId> recordIdSlot, boost::optional<value::SlotId> snapshotIdSlot, - boost::optional<value::SlotId> indexIdSlot, IndexKeysInclusionSet indexKeysToInclude, value::SlotVector vars, boost::optional<value::SlotId> seekKeySlotLow, @@ -128,7 +125,6 @@ private: const boost::optional<value::SlotId> _recordSlot; const boost::optional<value::SlotId> _recordIdSlot; const boost::optional<value::SlotId> _snapshotIdSlot; - const boost::optional<value::SlotId> _indexIdSlot; const IndexKeysInclusionSet _indexKeysToInclude; const value::SlotVector _vars; const boost::optional<value::SlotId> _seekKeySlotLow; @@ -145,14 +141,6 @@ private: std::unique_ptr<value::OwnedValueAccessor> _recordIdAccessor; std::unique_ptr<value::OwnedValueAccessor> _snapshotIdAccessor; - value::OwnedValueAccessor _indexIdAccessor; - value::ViewOfValueAccessor _indexIdViewAccessor; - - // This field holds the latest snapshot ID that we've received from _opCtx->recoveryUnit(). - // This field gets initialized by prepare(), and it gets updated each time doRestoreState() is - // called. - uint64_t _latestSnapshotId{0}; - // One accessor and slot for each key component that this stage will bind from an index entry's // KeyString. The accessors are in the same order as the key components they bind to. std::vector<value::OwnedValueAccessor> _accessors; diff --git a/src/mongo/db/exec/sbe/stages/plan_stats.h b/src/mongo/db/exec/sbe/stages/plan_stats.h index f487c7a45d2..263a4f86660 100644 --- a/src/mongo/db/exec/sbe/stages/plan_stats.h +++ b/src/mongo/db/exec/sbe/stages/plan_stats.h @@ -278,13 +278,8 @@ struct HashAggStats : public SpecificStats { } bool usedDisk{false}; - // The number of times that the entire hash table was spilled. - long long numSpills{0}; - // The number of individual records spilled to disk. long long spilledRecords{0}; - // An estimate, in bytes, of the size of the final spill table after all spill events have taken - // place. - long long spilledDataStorageSize{0}; + long long lastSpilledRecordSize{0}; }; struct HashLookupStats : public SpecificStats { diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp index 3d601d8779e..678d3f84ef9 100644 --- a/src/mongo/db/exec/sbe/stages/scan.cpp +++ b/src/mongo/db/exec/sbe/stages/scan.cpp @@ -209,7 +209,6 @@ void ScanStage::doSaveState(bool relinquishCursor) { cursor->setSaveStorageCursorOnDetachFromOperationContext(!relinquishCursor); } - _indexCatalogEntryMap.clear(); _coll.reset(); } @@ -392,14 +391,9 @@ PlanState ScanStage::getNext() { } // Return EOF if the index key is found to be inconsistent. - if (_scanCallbacks.indexKeyConsistencyCheckCallback && - !_scanCallbacks.indexKeyConsistencyCheckCallback(_opCtx, - _indexCatalogEntryMap, - _snapshotIdAccessor, - _indexIdAccessor, - _indexKeyAccessor, - _coll, - *nextRecord)) { + if (_scanCallbacks.indexKeyConsistencyCheckCallBack && + !_scanCallbacks.indexKeyConsistencyCheckCallBack( + _opCtx, _snapshotIdAccessor, _indexIdAccessor, _indexKeyAccessor, _coll, *nextRecord)) { return trackPlanState(PlanState::IS_EOF); } @@ -463,7 +457,6 @@ void ScanStage::close() { auto optTimer(getOptTimer(_opCtx)); trackClose(); - _indexCatalogEntryMap.clear(); _cursor.reset(); _randomCursor.reset(); _coll.reset(); @@ -749,7 +742,6 @@ void ParallelScanStage::doSaveState(bool relinquishCursor) { _cursor->save(); } - _indexCatalogEntryMap.clear(); _coll.reset(); } @@ -907,9 +899,8 @@ PlanState ParallelScanStage::getNext() { } // Return EOF if the index key is found to be inconsistent. - if (_scanCallbacks.indexKeyConsistencyCheckCallback && - !_scanCallbacks.indexKeyConsistencyCheckCallback(_opCtx, - _indexCatalogEntryMap, + if (_scanCallbacks.indexKeyConsistencyCheckCallBack && + !_scanCallbacks.indexKeyConsistencyCheckCallBack(_opCtx, _snapshotIdAccessor, _indexIdAccessor, _indexKeyAccessor, @@ -965,7 +956,6 @@ void ParallelScanStage::close() { auto optTimer(getOptTimer(_opCtx)); trackClose(); - _indexCatalogEntryMap.clear(); _cursor.reset(); _coll.reset(); _open = false; diff --git a/src/mongo/db/exec/sbe/stages/scan.h b/src/mongo/db/exec/sbe/stages/scan.h index 6f4980f52e1..37462ac5e14 100644 --- a/src/mongo/db/exec/sbe/stages/scan.h +++ b/src/mongo/db/exec/sbe/stages/scan.h @@ -45,11 +45,11 @@ struct ScanCallbacks { IndexKeyConsistencyCheckCallback indexKeyConsistencyCheck = {}, ScanOpenCallback scanOpen = {}) : indexKeyCorruptionCheckCallback(std::move(indexKeyCorruptionCheck)), - indexKeyConsistencyCheckCallback(std::move(indexKeyConsistencyCheck)), + indexKeyConsistencyCheckCallBack(std::move(indexKeyConsistencyCheck)), scanOpenCallback(std::move(scanOpen)) {} IndexKeyCorruptionCheckCallback indexKeyCorruptionCheckCallback; - IndexKeyConsistencyCheckCallback indexKeyConsistencyCheckCallback; + IndexKeyConsistencyCheckCallback indexKeyConsistencyCheckCallBack; ScanOpenCallback scanOpenCallback; }; @@ -83,12 +83,12 @@ struct ScanCallbacks { * * Debug string representations: * - * scan recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? indexKeySlot? - * indexKeyPatternSlot? [slot1 = fieldName1, ... slot_n = fieldName_n] collectionUuid + * scan recordSlot|none recordIdSlot|none snapshotIdSlot|none indexIdSlot|none indexKeySlot|none + * indexKeyPatternSlot|none [slot1 = fieldName1, ... slot_n = fieldName_n] collectionUuid * forward needOplogSlotForTs * - * seek seekKeySlot recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? indexKeySlot? - * indexKeyPatternSlot? [slot1 = fieldName1, ... slot_n = fieldName_n] + * seek seekKeySlot recordSlot|none recordIdSlot|none snapshotIdSlot|none indexIdSlot|none + * indexKeySlot|none indexKeyPatternSlot|none [slot1 = fieldName1, ... slot_n = fieldName_n] * collectionUuid forward needOplogSlotForTs */ class ScanStage final : public PlanStage { @@ -197,8 +197,6 @@ private: // collection is still valid. Only relevant to capped collections. bool _needsToCheckCappedPositionLost = false; - StringMap<const IndexCatalogEntry*> _indexCatalogEntryMap; - #if defined(MONGO_CONFIG_DEBUG_BUILD) // Debug-only buffer used to track the last thing returned from the stage. Between // saves/restores this is used to check that the storage cursor has not changed position. @@ -313,8 +311,6 @@ private: std::unique_ptr<SeekableRecordCursor> _cursor; - StringMap<const IndexCatalogEntry*> _indexCatalogEntryMap; - #if defined(MONGO_CONFIG_DEBUG_BUILD) // Debug-only buffer used to track the last thing returned from the stage. Between // saves/restores this is used to check that the storage cursor has not changed position. diff --git a/src/mongo/db/exec/sbe/stages/sort.cpp b/src/mongo/db/exec/sbe/stages/sort.cpp index eebdc0c2443..5acf73afe8d 100644 --- a/src/mongo/db/exec/sbe/stages/sort.cpp +++ b/src/mongo/db/exec/sbe/stages/sort.cpp @@ -213,11 +213,11 @@ void SortStage::open(bool reOpen) { _specificStats.totalDataSizeBytes += _sorter->totalDataSizeSorted(); _mergeIt.reset(_sorter->done()); - _specificStats.spills += _sorter->stats().spilledRanges(); + _specificStats.spills += _sorter->numSpills(); _specificStats.keysSorted += _sorter->numSorted(); auto& metricsCollector = ResourceConsumption::MetricsCollector::get(_opCtx); metricsCollector.incrementKeysSorted(_sorter->numSorted()); - metricsCollector.incrementSorterSpills(_sorter->stats().spilledRanges()); + metricsCollector.incrementSorterSpills(_sorter->numSpills()); _children[0]->close(); } diff --git a/src/mongo/db/exec/sbe/stages/union.cpp b/src/mongo/db/exec/sbe/stages/union.cpp index 3ddadb8c912..a661e6c579f 100644 --- a/src/mongo/db/exec/sbe/stages/union.cpp +++ b/src/mongo/db/exec/sbe/stages/union.cpp @@ -62,30 +62,28 @@ std::unique_ptr<PlanStage> UnionStage::clone() const { } void UnionStage::prepare(CompileCtx& ctx) { + value::SlotSet dupCheck; + for (size_t childNum = 0; childNum < _children.size(); childNum++) { _children[childNum]->prepare(ctx); } - // All of the slots listed in '_outputVals' must be unique. - value::SlotSet dupCheck; - for (auto slot : _outputVals) { - auto [it, inserted] = dupCheck.insert(slot); - uassert(4822807, str::stream() << "duplicate field: " << slot, inserted); - } - for (size_t idx = 0; idx < _outputVals.size(); ++idx) { std::vector<value::SlotAccessor*> accessors; accessors.reserve(_children.size()); for (size_t childNum = 0; childNum < _children.size(); childNum++) { - // Slots listed in '_inputVals' may not appear in '_outputVals'. auto slot = _inputVals[childNum][idx]; - bool slotFound = dupCheck.count(slot); - uassert(4822806, str::stream() << "duplicate field: " << slot, !slotFound); + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822806, str::stream() << "duplicate field: " << slot, inserted); accessors.emplace_back(_children[childNum]->getAccessor(ctx, slot)); } + auto slot = _outputVals[idx]; + auto [it, inserted] = dupCheck.insert(slot); + uassert(4822807, str::stream() << "duplicate field: " << slot, inserted); + _outValueAccessors.emplace_back(value::SwitchAccessor{std::move(accessors)}); } } diff --git a/src/mongo/db/exec/sbe/util/spilling.cpp b/src/mongo/db/exec/sbe/util/spilling.cpp index 7675fad6846..c54f3bfe956 100644 --- a/src/mongo/db/exec/sbe/util/spilling.cpp +++ b/src/mongo/db/exec/sbe/util/spilling.cpp @@ -29,18 +29,6 @@ #include "mongo/db/exec/sbe/util/spilling.h" -#include "mongo/base/status.h" -#include "mongo/base/status_with.h" -#include "mongo/base/string_data.h" -#include "mongo/bson/timestamp.h" -#include "mongo/db/query/query_knobs_gen.h" -#include "mongo/db/storage/record_data.h" -#include "mongo/db/storage/recovery_unit.h" -#include "mongo/db/storage/write_unit_of_work.h" -#include "mongo/util/assert_util.h" -#include "mongo/util/bufreader.h" -#include "mongo/util/str.h" - namespace mongo { namespace sbe { @@ -69,136 +57,56 @@ KeyString::Value decodeKeyString(const RecordId& rid, KeyString::TypeBits typeBi return kb.getValueCopy(); } -SpillingStore::SpillingStore(OperationContext* opCtx, KeyFormat format) { - _recordStore = - opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(opCtx, format); - - _spillingUnit = std::unique_ptr<RecoveryUnit>( - opCtx->getServiceContext()->getStorageEngine()->newRecoveryUnit()); - _spillingUnit->setCacheMaxWaitTimeout(Milliseconds(internalQuerySpillingMaxWaitTimeout.load())); - _spillingState = WriteUnitOfWork::RecoveryUnitState::kNotInUnitOfWork; -} - -SpillingStore::~SpillingStore() {} - -int SpillingStore::upsertToRecordStore(OperationContext* opCtx, - const RecordId& recordKey, - const value::MaterializedRow& key, - const value::MaterializedRow& val, - bool update) { - BufBuilder buf; - key.serializeForSorter(buf); - val.serializeForSorter(buf); - return upsertToRecordStore(opCtx, recordKey, buf, update); +boost::optional<value::MaterializedRow> readFromRecordStore(OperationContext* opCtx, + RecordStore* rs, + const RecordId& rid) { + RecordData record; + if (rs->findRecord(opCtx, rid, &record)) { + auto valueReader = BufReader(record.data(), record.size()); + return value::MaterializedRow::deserializeForSorter(valueReader, {}); + } + return boost::none; } -int SpillingStore::upsertToRecordStore( - OperationContext* opCtx, - const RecordId& key, - const value::MaterializedRow& val, - const KeyString::TypeBits& typeBits, // recover type of value. - bool update) { +int upsertToRecordStore(OperationContext* opCtx, + RecordStore* rs, + const RecordId& key, + const value::MaterializedRow& val, + const KeyString::TypeBits& typeBits, // recover type of value. + bool update) { BufBuilder bufValue; val.serializeForSorter(bufValue); - // Append the 'typeBits' to the end of the val's buffer so the 'key' can be reconstructed when - // draining HashAgg. - bufValue.appendBuf(typeBits.getBuffer(), typeBits.getSize()); - - return upsertToRecordStore(opCtx, key, bufValue, update); + return upsertToRecordStore(opCtx, rs, key, bufValue, typeBits, update); } -int SpillingStore::upsertToRecordStore( - OperationContext* opCtx, - const RecordId& key, - BufBuilder& buf, - const KeyString::TypeBits& typeBits, // recover type of value. - bool update) { +int upsertToRecordStore(OperationContext* opCtx, + RecordStore* rs, + const RecordId& key, + BufBuilder& buf, + const KeyString::TypeBits& typeBits, // recover type of value. + bool update) { + // Append the 'typeBits' to the end of the val's buffer so the 'key' can be reconstructed when // draining HashAgg. buf.appendBuf(typeBits.getBuffer(), typeBits.getSize()); - return upsertToRecordStore(opCtx, key, buf, update); -} - -int SpillingStore::upsertToRecordStore(OperationContext* opCtx, - const RecordId& key, - BufBuilder& buf, - bool update) { assertIgnorePrepareConflictsBehavior(opCtx); - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); WriteUnitOfWork wuow(opCtx); auto result = mongo::Status::OK(); if (update) { - result = rs()->updateRecord(opCtx, key, buf.buf(), buf.len()); + result = rs->updateRecord(opCtx, key, buf.buf(), buf.len()); } else { - auto status = rs()->insertRecord(opCtx, key, buf.buf(), buf.len(), Timestamp{}); + auto status = rs->insertRecord(opCtx, key, buf.buf(), buf.len(), Timestamp{}); result = status.getStatus(); } wuow.commit(); - if (!result.isOK()) { tasserted(5843600, str::stream() << "Failed to write to disk because " << result.reason()); return 0; } return buf.len(); } - -Status SpillingStore::insertRecords(OperationContext* opCtx, - std::vector<Record>* inOutRecords, - const std::vector<Timestamp>& timestamps) { - assertIgnorePrepareConflictsBehavior(opCtx); - - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); - WriteUnitOfWork wuow(opCtx); - auto status = rs()->insertRecords(opCtx, inOutRecords, timestamps); - wuow.commit(); - - return status; -} - -boost::optional<value::MaterializedRow> SpillingStore::readFromRecordStore(OperationContext* opCtx, - const RecordId& rid) { - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); - - RecordData record; - if (rs()->findRecord(opCtx, rid, &record)) { - auto valueReader = BufReader(record.data(), record.size()); - return value::MaterializedRow::deserializeForSorter(valueReader, {}); - } - return boost::none; -} - -bool SpillingStore::findRecord(OperationContext* opCtx, const RecordId& loc, RecordData* out) { - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); - - return rs()->findRecord(opCtx, loc, out); -} - -void SpillingStore::switchToSpilling(OperationContext* opCtx) { - invariant(!_originalUnit); - _originalUnit = opCtx->releaseRecoveryUnit(); - _originalState = opCtx->setRecoveryUnit(std::move(_spillingUnit), _spillingState); -} -void SpillingStore::switchToOriginal(OperationContext* opCtx) { - invariant(!_spillingUnit); - _spillingUnit = opCtx->releaseRecoveryUnit(); - _spillingState = opCtx->setRecoveryUnit(std::move(_originalUnit), _originalState); - invariant(!(_spillingUnit->getState() == RecoveryUnit::State::kInactiveInUnitOfWork || - _spillingUnit->getState() == RecoveryUnit::State::kActive)); -} - -void SpillingStore::saveState() { - _spillingUnit->abandonSnapshot(); -} -void SpillingStore::restoreState() { - // We do not have to do anything. -} - } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/util/spilling.h b/src/mongo/db/exec/sbe/util/spilling.h index 205d6f1a031..95f73e2b02e 100644 --- a/src/mongo/db/exec/sbe/util/spilling.h +++ b/src/mongo/db/exec/sbe/util/spilling.h @@ -29,14 +29,9 @@ #pragma once -#include <boost/optional/optional.hpp> -#include <utility> +#include "mongo/platform/basic.h" -#include "mongo/bson/util/builder.h" #include "mongo/db/exec/sbe/values/slot.h" -#include "mongo/db/operation_context.h" -#include "mongo/db/record_id.h" -#include "mongo/db/storage/record_store.h" #include "mongo/db/storage/temporary_record_store.h" namespace mongo { @@ -55,104 +50,27 @@ std::pair<RecordId, KeyString::TypeBits> encodeKeyString(KeyString::Builder&, // Reconstructs the KeyString carried in RecordId using 'typeBits'. KeyString::Value decodeKeyString(const RecordId& rid, KeyString::TypeBits typeBits); -/** - * SpillingStore is a wrapper around a temporary record store than maintains its own transaction as - * we do not want to intermingle operations running in the main query with spill reads and writes. - */ -class SpillingStore { -public: - SpillingStore(OperationContext* opCtx, KeyFormat format = KeyFormat::String); - ~SpillingStore(); - - /** - * When a collator is provided, the key is encoded using the collator before being converted to - * a record id. In this case, it is not possible to recover the key from the record id, thus we - * need to store the original value of the key as well. - */ - int upsertToRecordStore(OperationContext* opCtx, - const RecordId& recordKey, - const value::MaterializedRow& key, - const value::MaterializedRow& val, - bool update); - /** - * Inserts or updates a key/value into 'rs'. The 'update' flag controls whether or not an update - * will be performed. If a key/value pair is inserted into the 'rs' that already exists and - * 'update' is false, this function will tassert. - * - * Returns the size of the new record in bytes, including the record id and value portions. - */ - int upsertToRecordStore(OperationContext* opCtx, - const RecordId& key, - const value::MaterializedRow& val, - const KeyString::TypeBits& typeBits, - bool update); - int upsertToRecordStore(OperationContext* opCtx, - const RecordId& key, - BufBuilder& buf, - const KeyString::TypeBits& typeBits, // recover type of value. - bool update); - int upsertToRecordStore(OperationContext* opCtx, - const RecordId& key, - BufBuilder& buf, - bool update); - - - Status insertRecords(OperationContext* opCtx, - std::vector<Record>* inOutRecords, - const std::vector<Timestamp>& timestamps); - - // Reads a materialized row from the record store. - boost::optional<value::MaterializedRow> readFromRecordStore(OperationContext* opCtx, - const RecordId& rid); - - bool findRecord(OperationContext* opCtx, const RecordId& loc, RecordData* out); - - auto rs() { - return _recordStore->rs(); - } - - auto getCursor(OperationContext* opCtx) { - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); - return rs()->getCursor(opCtx); - } - - void resetCursor(OperationContext* opCtx, std::unique_ptr<SeekableRecordCursor>& cursor) { - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); - cursor.reset(); - } +// Reads a materialized row from the record store. +boost::optional<value::MaterializedRow> readFromRecordStore(OperationContext* opCtx, + RecordStore* rs, + const RecordId& rid); - auto saveCursor(OperationContext* opCtx, std::unique_ptr<SeekableRecordCursor>& cursor) { - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); - - return cursor->save(); - } - - auto restoreCursor(OperationContext* opCtx, std::unique_ptr<SeekableRecordCursor>& cursor) { - switchToSpilling(opCtx); - ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); }); - - return cursor->restore(); - } - - void saveState(); - void restoreState(); - -private: - void switchToSpilling(OperationContext* opCtx); - void switchToOriginal(OperationContext* opCtx); - - std::unique_ptr<TemporaryRecordStore> _recordStore; - - std::unique_ptr<RecoveryUnit> _originalUnit; - WriteUnitOfWork::RecoveryUnitState _originalState; - - std::unique_ptr<RecoveryUnit> _spillingUnit; - WriteUnitOfWork::RecoveryUnitState _spillingState; - - size_t _counter{0}; -}; +/** Inserts or updates a key/value into 'rs'. The 'update' flag controls whether or not an update + * will be performed. If a key/value pair is inserted into the 'rs' that already exists and + * 'update' is false, this function will tassert. + */ +int upsertToRecordStore(OperationContext* opCtx, + RecordStore* rs, + const RecordId& key, + const value::MaterializedRow& val, + const KeyString::TypeBits& typeBits, + bool update); + +int upsertToRecordStore(OperationContext* opCtx, + RecordStore* rs, + const RecordId& key, + BufBuilder& buf, + const KeyString::TypeBits& typeBits, // recover type of value. + bool update); } // namespace sbe } // namespace mongo diff --git a/src/mongo/db/exec/sbe/values/slot.cpp b/src/mongo/db/exec/sbe/values/slot.cpp index 93605a9abfa..45cbc977980 100644 --- a/src/mongo/db/exec/sbe/values/slot.cpp +++ b/src/mongo/db/exec/sbe/values/slot.cpp @@ -42,9 +42,7 @@ #include "mongo/util/bufreader.h" namespace mongo::sbe::value { - -static std::pair<TypeTags, Value> deserializeValue(BufReader& buf, - const CollatorInterface* collator) { +static std::pair<TypeTags, Value> deserializeValue(BufReader& buf) { auto tag = static_cast<TypeTags>(buf.read<uint8_t>()); Value val; @@ -109,7 +107,7 @@ static std::pair<TypeTags, Value> deserializeValue(BufReader& buf, if (cnt) { arr->reserve(cnt); for (size_t idx = 0; idx < cnt; ++idx) { - auto [tag, val] = deserializeValue(buf, collator); + auto [tag, val] = deserializeValue(buf); arr->push_back(tag, val); } } @@ -118,16 +116,13 @@ static std::pair<TypeTags, Value> deserializeValue(BufReader& buf, break; } case TypeTags::ArraySet: { - // The first byte is a flag to tell us whether the ArraySet had a collation prior to - // serialization. - auto collated = buf.read<char>(); - auto [arrTag, arrVal] = makeNewArraySet(collated ? collator : nullptr); auto cnt = buf.read<LittleEndian<size_t>>(); + auto [arrTag, arrVal] = makeNewArraySet(); auto arr = getArraySetView(arrVal); if (cnt) { arr->reserve(cnt); for (size_t idx = 0; idx < cnt; ++idx) { - auto [tag, val] = deserializeValue(buf, collator); + auto [tag, val] = deserializeValue(buf); arr->push_back(tag, val); } } @@ -143,7 +138,7 @@ static std::pair<TypeTags, Value> deserializeValue(BufReader& buf, obj->reserve(cnt); for (size_t idx = 0; idx < cnt; ++idx) { auto fieldName = buf.readCStr(); - auto [tag, val] = deserializeValue(buf, collator); + auto [tag, val] = deserializeValue(buf); obj->push_back(fieldName, tag, val); } } @@ -218,12 +213,12 @@ static std::pair<TypeTags, Value> deserializeValue(BufReader& buf, } MaterializedRow MaterializedRow::deserializeForSorter(BufReader& buf, - const SorterDeserializeSettings& settings) { + const SorterDeserializeSettings&) { auto cnt = buf.read<LittleEndian<size_t>>(); MaterializedRow result{cnt}; for (size_t idx = 0; idx < cnt; ++idx) { - auto [tag, val] = deserializeValue(buf, settings.collator); + auto [tag, val] = deserializeValue(buf); result.reset(idx, true, tag, val); } @@ -293,12 +288,6 @@ static void serializeValue(BufBuilder& buf, TypeTags tag, Value val) { } case TypeTags::ArraySet: { auto arr = getArraySetView(val); - // If an ArraySet has a collation, we serialize a byte which acts as a flag as to - // whether the set should be created with a collation upon deserialization. Also, we - // assume that the caller which does deserialization will have the context about what - // the collation is, and therefore we can save space by not serializing the full - // description of the collation. - buf.appendChar(arr->getCollator() ? 1 : 0); buf.appendNum(arr->size()); for (auto& kv : arr->values()) { serializeValue(buf, kv.first, kv.second); @@ -382,22 +371,7 @@ static void serializeValue(BufBuilder& buf, TypeTags tag, Value val) { } } -/** - * If non-null 'collator' is provided during serialization, then the encoding guarantees that values - * which are equal up to the collation will encode to the same result, allowing for collation-aware - * equality comparisons. However, the collator-aware encoded values are not always decodable. Groups - * store an original copy of the group key alongside the encoded key string when there is a - * collation, so the key string does not need to be decodable. - */ -static void serializeValueIntoKeyString(KeyString::Builder& buf, - TypeTags tag, - Value val, - const CollatorInterface* collator) { - - const auto stringTransformFn = [&](StringData stringData) { - return collator->getComparisonString(stringData); - }; - +static void serializeValueIntoKeyString(KeyString::Builder& buf, TypeTags tag, Value val) { switch (tag) { case TypeTags::Nothing: { buf.appendBool(false); @@ -460,21 +434,19 @@ static void serializeValueIntoKeyString(KeyString::Builder& buf, buf.appendUndefined(); break; } - case TypeTags::StringSmall: + case TypeTags::StringSmall: { + // Small strings cannot contain null bytes, so it is safe to serialize them as plain + // C-strings with a null terminator. + buf.appendBool(true); + buf.appendString(getStringView(tag, val)); + break; + } case TypeTags::StringBig: case TypeTags::bsonString: { buf.appendBool(true); - if (collator) { - buf.appendString(getStringView(tag, val), stringTransformFn); - } else { - buf.appendString(getStringView(tag, val)); - } + buf.appendString(getStringOrSymbolView(tag, val)); break; } - // Note that the collation would have to apply to bsonSymbol values to match the - // behavior of the Classic engine when spilling groups to disk. bsonSymbol is - // deprecated, however, and SBE no longer provides strict correctness guarantees about - // computations on bsonSymbol values. case TypeTags::bsonSymbol: { buf.appendBool(true); buf.appendSymbol(getStringOrSymbolView(tag, val)); @@ -485,13 +457,9 @@ static void serializeValueIntoKeyString(KeyString::Builder& buf, // TODO SERVER-61629: convert this to serialize the 'arr' directly instead of // constructing a BSONArray. BSONArrayBuilder builder; - bson::convertToBsonObj(builder, value::ArrayEnumerator{tag, val}); + bson::convertToBsonObj(builder, getArrayView(val)); buf.appendBool(true); - if (collator) { - buf.appendArray(builder.arr(), stringTransformFn); - } else { - buf.appendArray(builder.arr()); - } + buf.appendArray(BSONArray(builder.done())); break; } case TypeTags::Object: { @@ -500,37 +468,27 @@ static void serializeValueIntoKeyString(KeyString::Builder& buf, BSONObjBuilder builder; bson::convertToBsonObj(builder, getObjectView(val)); buf.appendBool(true); - if (collator) { - buf.appendObject(builder.done(), stringTransformFn); - } else { - buf.appendObject(builder.done()); - } + buf.appendObject(builder.done()); break; } - case TypeTags::bsonObjectId: case TypeTags::ObjectId: { buf.appendBool(true); - buf.appendOID(OID::from(getRawPointerView(val))); + buf.appendBytes(getObjectIdView(val), sizeof(ObjectIdType)); break; } case TypeTags::bsonObject: { - BSONObj bson{getRawPointerView(val)}; buf.appendBool(true); - if (collator) { - buf.appendObject(bson, stringTransformFn); - } else { - buf.appendObject(bson); - } + buf.appendObject(BSONObj(getRawPointerView(val))); break; } case TypeTags::bsonArray: { - BSONObj bson{getRawPointerView(val)}; buf.appendBool(true); - if (collator) { - buf.appendArray(BSONArray(BSONObj(bson)), stringTransformFn); - } else { - buf.appendArray(BSONArray(BSONObj(bson))); - } + buf.appendArray(BSONArray(BSONObj(getRawPointerView(val)))); + break; + } + case TypeTags::bsonObjectId: { + buf.appendBool(true); + buf.appendOID(OID::from(getRawPointerView(val))); break; } case TypeTags::bsonBinData: { @@ -599,18 +557,15 @@ void MaterializedRow::serializeForSorter(BufBuilder& buf) const { } } -void MaterializedRow::serializeIntoKeyString(KeyString::Builder& buf, - const CollatorInterface* collator) const { +void MaterializedRow::serializeIntoKeyString(KeyString::Builder& buf) const { for (size_t idx = 0; idx < size(); ++idx) { auto [tag, val] = getViewOfValue(idx); - serializeValueIntoKeyString(buf, tag, val, collator); + serializeValueIntoKeyString(buf, tag, val); } } -MaterializedRow MaterializedRow::deserializeFromKeyString( - const KeyString::Value& keyString, - BufBuilder* valueBufferBuilder, - boost::optional<size_t> numPrefixValsToRead) { +MaterializedRow MaterializedRow::deserializeFromKeyString(const KeyString::Value& keyString, + BufBuilder* valueBufferBuilder) { BufReader reader(keyString.getBuffer(), keyString.getSize()); KeyString::TypeBits typeBits(keyString.getTypeBits()); KeyString::TypeBits::Reader typeBitsReader(typeBits); @@ -622,8 +577,7 @@ MaterializedRow MaterializedRow::deserializeFromKeyString( &reader, &typeBitsReader, false /* inverted */, typeBits.version, &valBuilder); } while (keepReading); - size_t sizeOfRow = numPrefixValsToRead ? *numPrefixValsToRead : valBuilder.numValues(); - MaterializedRow result{sizeOfRow}; + MaterializedRow result{valBuilder.numValues()}; valBuilder.readValues(result); return result; diff --git a/src/mongo/db/exec/sbe/values/slot.h b/src/mongo/db/exec/sbe/values/slot.h index e682b2bfeb8..f853f816d4d 100644 --- a/src/mongo/db/exec/sbe/values/slot.h +++ b/src/mongo/db/exec/sbe/values/slot.h @@ -468,11 +468,8 @@ public: } // The following methods are used by the sorter only. - struct SorterDeserializeSettings { - const CollatorInterface* collator{nullptr}; - }; - static MaterializedRow deserializeForSorter(BufReader& buf, - const SorterDeserializeSettings& settings); + struct SorterDeserializeSettings {}; + static MaterializedRow deserializeForSorter(BufReader& buf, const SorterDeserializeSettings&); void serializeForSorter(BufBuilder& buf) const; int memUsageForSorter() const; auto getOwned() const { @@ -486,21 +483,11 @@ public: * intended for spilling key values used in the HashAgg stage. The format is not guaranteed to * be stable between versions, so it should not be used for long-term storage or communication * between instances. - * - * If 'numPrefixValsToRead' is provided, then only the given number of values from 'keyString' - * are decoded into the resulting 'MaterializedRow'. The remaining suffix values in the - * 'keyString' are ignored. - * - * If non-null 'collator' is provided during serialization, then any strings in the row are - * encoded as ICU collation keys prior to being KeyString-encoded. */ - static MaterializedRow deserializeFromKeyString( - const KeyString::Value& keyString, - BufBuilder* valueBufferBuilder, - boost::optional<size_t> numPrefixValsToRead = boost::none); + static MaterializedRow deserializeFromKeyString(const KeyString::Value& keyString, - void serializeIntoKeyString(KeyString::Builder& builder, - const CollatorInterface* collator = nullptr) const; + BufBuilder* valueBufferBuilder); + void serializeIntoKeyString(KeyString::Builder& builder) const; private: static size_t sizeInBytes(size_t count) { @@ -590,9 +577,9 @@ private: }; struct MaterializedRowEq { - using ComparatorType = StringData::ComparatorInterface; + using ComparatorType = StringData::ComparatorInterface*; - explicit MaterializedRowEq(const ComparatorType* comparator = nullptr) + explicit MaterializedRowEq(const ComparatorType comparator = nullptr) : _comparator(comparator) {} bool operator()(const MaterializedRow& lhs, const MaterializedRow& rhs) const { @@ -610,7 +597,7 @@ struct MaterializedRowEq { } private: - const ComparatorType* _comparator = nullptr; + const ComparatorType _comparator = nullptr; }; struct MaterializedRowLess { diff --git a/src/mongo/db/exec/sbe/values/value.cpp b/src/mongo/db/exec/sbe/values/value.cpp index e0f73c87ddf..6c21f4f6e5d 100644 --- a/src/mongo/db/exec/sbe/values/value.cpp +++ b/src/mongo/db/exec/sbe/values/value.cpp @@ -860,7 +860,7 @@ bool isInfinity(TypeTags tag, Value val) { (tag == TypeTags::NumberDecimal && bitcastTo<Decimal128>(val).isInfinite()); } -bool ArraySet::push_back(TypeTags tag, Value val) { +void ArraySet::push_back(TypeTags tag, Value val) { if (tag != TypeTags::Nothing) { ValueGuard guard{tag, val}; auto [it, inserted] = _values.insert({tag, val}); @@ -868,11 +868,7 @@ bool ArraySet::push_back(TypeTags tag, Value val) { if (inserted) { guard.reset(); } - - return inserted; } - - return false; } std::pair<TypeTags, Value> ArrayEnumerator::getViewOfValue() const { diff --git a/src/mongo/db/exec/sbe/values/value.h b/src/mongo/db/exec/sbe/values/value.h index ae207106d8b..26ba3e5e1de 100644 --- a/src/mongo/db/exec/sbe/values/value.h +++ b/src/mongo/db/exec/sbe/values/value.h @@ -844,14 +844,7 @@ public: } } - /** - * Adds the given SBE value to the set if an equal value is not already present. Assumes - * ownership of the given value. - * - * Returns true if the value was newly inserted, otherwise returns false to indicate that an - * equal value was already present in the set. - */ - bool push_back(TypeTags tag, Value val); + void push_back(TypeTags tag, Value val); auto& values() const noexcept { return _values; diff --git a/src/mongo/db/exec/sbe/values/value_builder.h b/src/mongo/db/exec/sbe/values/value_builder.h index 00333e9f824..9ad2b511242 100644 --- a/src/mongo/db/exec/sbe/values/value_builder.h +++ b/src/mongo/db/exec/sbe/values/value_builder.h @@ -191,11 +191,8 @@ public: virtual size_t numValues() const = 0; protected: - // We expect most rows to end up containing this many values or fewer. - static constexpr int kInlinedVectorSize = 16; - std::pair<TypeTags, Value> getValue(size_t index, int bufferLen) { - invariant(index < _tagList.size()); + invariant(index < _numValues); auto tag = _tagList[index]; auto val = _valList[index]; @@ -227,8 +224,9 @@ protected: } void appendValue(TypeTags tag, Value val) noexcept { - _tagList.push_back(tag); - _valList.push_back(val); + _tagList[_numValues] = tag; + _valList[_numValues] = val; + ++_numValues; } void appendValue(std::pair<TypeTags, Value> in) noexcept { @@ -243,12 +241,14 @@ protected: // storing a pointer, we store an _offset_ into the under-construction buffer. Translation from // offset to pointer occurs as part of the 'readValues()' function. void appendValueBufferOffset(TypeTags tag) { - _tagList.push_back(tag); - _valList.push_back(value::bitcastFrom<int32_t>(_valueBufferBuilder->len())); + _tagList[_numValues] = tag; + _valList[_numValues] = value::bitcastFrom<int32_t>(_valueBufferBuilder->len()); + ++_numValues; } - absl::InlinedVector<TypeTags, kInlinedVectorSize> _tagList; - absl::InlinedVector<Value, kInlinedVectorSize> _valList; + std::array<TypeTags, Ordering::kMaxCompoundIndexKeys> _tagList; + std::array<Value, Ordering::kMaxCompoundIndexKeys> _valList; + size_t _numValues = 0; BufBuilder* _valueBufferBuilder; }; @@ -270,12 +270,11 @@ public: // buffer, this value will remain in that buffer, even though we've removed it from the // list. It will still get deallocated along with everything else when that buffer gets // cleared or deleted, though, so there is no leak. - _tagList.pop_back(); - _valList.pop_back(); + --_numValues; } size_t numValues() const override { - return _tagList.size(); + return _numValues; } /** @@ -285,7 +284,7 @@ public: */ void readValues(std::vector<OwnedValueAccessor>* accessors) { auto bufferLen = _valueBufferBuilder->len(); - for (size_t i = 0; i < _tagList.size(); ++i) { + for (size_t i = 0; i < _numValues; ++i) { auto [tag, val] = getValue(i, bufferLen); invariant(i < accessors->size()); (*accessors)[i].reset(false, tag, val); @@ -305,7 +304,7 @@ public: size_t numValues() const override { size_t nVals = 0; size_t bufIdx = 0; - while (bufIdx < _tagList.size()) { + while (bufIdx < _numValues) { auto tag = _tagList[bufIdx]; auto val = _valList[bufIdx]; if (tag == TypeTags::Boolean && !bitcastTo<bool>(val)) { @@ -324,10 +323,7 @@ public: auto bufferLen = _valueBufferBuilder->len(); size_t bufIdx = 0; size_t rowIdx = 0; - // The 'row' output parameter might be smaller than the number of values owned by this - // builder. Be careful to only read as many values into 'row' as this output 'row' has space - // for. - while (rowIdx < row.size()) { + while (bufIdx < _numValues) { invariant(rowIdx < row.size()); auto [_, tagNothing, valNothing] = getValue(bufIdx++, bufferLen); tassert(6136200, "sbe tag must be 'Boolean'", tagNothing == TypeTags::Boolean); diff --git a/src/mongo/db/exec/sbe/values/value_serialization_test.cpp b/src/mongo/db/exec/sbe/values/value_serialization_test.cpp index ceacf610a4a..5b44bef6549 100644 --- a/src/mongo/db/exec/sbe/values/value_serialization_test.cpp +++ b/src/mongo/db/exec/sbe/values/value_serialization_test.cpp @@ -268,18 +268,6 @@ TEST_F(ValueSerializeForKeyString, SbeArray) { runTest({{testDataTag, testDataVal}}); } -TEST_F(ValueSerializeForKeyString, ArraySet) { - auto [tag, val] = sbe::value::makeNewArraySet(); - sbe::value::ValueGuard guard{tag, val}; - auto* arraySet = sbe::value::getArraySetView(val); - - arraySet->push_back(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(1)); - arraySet->push_back(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(2)); - arraySet->push_back(value::TypeTags::NumberDouble, value::bitcastFrom<double>(3.0)); - - runTest({{tag, val}}); -} - TEST_F(ValueSerializeForKeyString, DateTime) { runTest({{value::TypeTags::Date, value::bitcastFrom<int64_t>(1234)}, {value::TypeTags::Timestamp, value::bitcastFrom<uint64_t>(5678)}}); @@ -454,30 +442,4 @@ TEST_F(ValueSerializeForKeyString, BsonCodeWScope) { runTest({{cwsTag1, cwsVal1}, {cwsTag2, cwsVal2}, {cwsTag3, cwsVal3}}); } - -// Test that roundtripping through KeyString works for a wide row. KeyStrings used in indexes are -// typically constrained in the number of components they can have, since we limit compound indexes -// to at most 32 components. But roundtripping rows wider than 32 still needs to work. -// -// This test was originally designed to reproduce SERVER-76321. -TEST_F(ValueSerializeForKeyString, RoundtripWideRow) { - std::vector<std::pair<sbe::value::TypeTags, sbe::value::Value>> row; - for (int32_t i = 0; i < 40; ++i) { - row.emplace_back(sbe::value::TypeTags::NumberInt32, sbe::value::bitcastFrom<int32_t>(i)); - } - runTest(row); -} - -// Test that roundtripping through KeyString works for ObjectIdType: ObjectId; bsonObjectId. -TEST_F(ValueSerializeForKeyString, RoundtripObjectIdType) { - auto [objectIdTag, objectIdVal] = value::makeNewObjectId(); - - auto oid = OID::gen(); - auto obj = BSON("" << oid); - auto oidStorage = obj.firstElement().value(); - - sbe::value::ValueGuard testDataGuard{objectIdTag, objectIdVal}; - runTest({{objectIdTag, objectIdVal}, - {value::TypeTags::bsonObjectId, value::bitcastFrom<const char*>(oidStorage)}}); -} } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/vm/arith.cpp b/src/mongo/db/exec/sbe/vm/arith.cpp index e6c9d380ffd..e4c67775ad8 100644 --- a/src/mongo/db/exec/sbe/vm/arith.cpp +++ b/src/mongo/db/exec/sbe/vm/arith.cpp @@ -503,57 +503,6 @@ void ByteCode::aggDoubleDoubleSumImpl(value::Array* arr, } } -void ByteCode::aggMergeDoubleDoubleSumsImpl(value::Array* accumulator, - value::TypeTags rhsTag, - value::Value rhsValue) { - auto [accumWidestType, _1] = accumulator->getAt(AggSumValueElems::kNonDecimalTotalTag); - - tassert(7039532, "value must be of type 'Array'", rhsTag == value::TypeTags::Array); - auto nextDoubleDoubleArr = value::getArrayView(rhsValue); - - tassert(7039533, - "array does not have enough elements", - nextDoubleDoubleArr->size() >= AggSumValueElems::kMaxSizeOfArray - 1); - - // First aggregate the non-decimal sum, then the non-decimal addend. Both should be doubles. - auto [sumTag, sum] = nextDoubleDoubleArr->getAt(AggSumValueElems::kNonDecimalTotalSum); - tassert(7039534, "expected 'NumberDouble'", sumTag == value::TypeTags::NumberDouble); - aggDoubleDoubleSumImpl(accumulator, sumTag, sum); - - auto [addendTag, addend] = nextDoubleDoubleArr->getAt(AggSumValueElems::kNonDecimalTotalAddend); - tassert(7039535, "expected 'NumberDouble'", addendTag == value::TypeTags::NumberDouble); - // There is a special case when the 'sum' is infinite and the 'addend' is NaN. This DoubleDouble - // value represents infinity, not NaN. Therefore, we avoid incorporating the NaN 'addend' value - // into the sum. - if (std::isfinite(value::bitcastTo<double>(sum)) || - !std::isnan(value::bitcastTo<double>(addend))) { - aggDoubleDoubleSumImpl(accumulator, addendTag, addend); - } - - // Determine the widest non-decimal type that we've seen so far, and set the accumulator state - // accordingly. We do this after computing the sums, since 'aggDoubleDoubleSumImpl()' will - // set the widest type to 'NumberDouble' when we call it above. - auto [newValWidestType, _2] = nextDoubleDoubleArr->getAt(AggSumValueElems::kNonDecimalTotalTag); - tassert( - 7039536, "unexpected 'NumberDecimal'", newValWidestType != value::TypeTags::NumberDecimal); - tassert( - 7039537, "unexpected 'NumberDecimal'", accumWidestType != value::TypeTags::NumberDecimal); - auto widestType = getWidestNumericalType(newValWidestType, accumWidestType); - accumulator->setAt( - AggSumValueElems::kNonDecimalTotalTag, widestType, value::bitcastFrom<int32_t>(0)); - - // If there's a decimal128 sum as part of the incoming DoubleDouble sum, incorporate it into the - // accumulator. - if (nextDoubleDoubleArr->size() == AggSumValueElems::kMaxSizeOfArray) { - auto [decimalTotalTag, decimalTotalVal] = - nextDoubleDoubleArr->getAt(AggSumValueElems::kDecimalTotal); - tassert(7039538, - "The decimalTotal must be 'NumberDecimal'", - decimalTotalTag == TypeTags::NumberDecimal); - aggDoubleDoubleSumImpl(accumulator, decimalTotalTag, decimalTotalVal); - } -} - void ByteCode::aggStdDevImpl(value::Array* arr, value::TypeTags rhsTag, value::Value rhsValue) { if (!isNumber(rhsTag)) { return; @@ -602,67 +551,6 @@ void ByteCode::aggStdDevImpl(value::Array* arr, value::TypeTags rhsTag, value::V return setStdDevArray(newCountVal, newMeanVal, newM2Val, arr); } -void ByteCode::aggMergeStdDevsImpl(value::Array* accumulator, - value::TypeTags rhsTag, - value::Value rhsValue) { - tassert(7039542, "expected value of type 'Array'", rhsTag == value::TypeTags::Array); - auto nextArr = value::getArrayView(rhsValue); - - tassert(7039543, - "expected array to have exactly 3 elements", - accumulator->size() == AggStdDevValueElems::kSizeOfArray); - tassert(7039544, - "expected array to have exactly 3 elements", - nextArr->size() == AggStdDevValueElems::kSizeOfArray); - - auto [newCountTag, newCountVal] = nextArr->getAt(AggStdDevValueElems::kCount); - tassert(7039545, "expected 64-bit int", newCountTag == value::TypeTags::NumberInt64); - int64_t newCount = value::bitcastTo<int64_t>(newCountVal); - - // If the incoming partial aggregate has a count of zero, then it represents the partial - // standard deviation of no data points. This means that it can be safely ignored, and we return - // the accumulator as is. - if (newCount == 0) { - return; - } - - auto [oldCountTag, oldCountVal] = accumulator->getAt(AggStdDevValueElems::kCount); - tassert(7039546, "expected 64-bit int", oldCountTag == value::TypeTags::NumberInt64); - int64_t oldCount = value::bitcastTo<int64_t>(oldCountVal); - - auto [oldMeanTag, oldMeanVal] = accumulator->getAt(AggStdDevValueElems::kRunningMean); - tassert(7039547, "expected double", oldMeanTag == value::TypeTags::NumberDouble); - double oldMean = value::bitcastTo<double>(oldMeanVal); - - auto [newMeanTag, newMeanVal] = nextArr->getAt(AggStdDevValueElems::kRunningMean); - tassert(7039548, "expected double", newMeanTag == value::TypeTags::NumberDouble); - double newMean = value::bitcastTo<double>(newMeanVal); - - auto [oldM2Tag, oldM2Val] = accumulator->getAt(AggStdDevValueElems::kRunningM2); - tassert(7039531, "expected double", oldM2Tag == value::TypeTags::NumberDouble); - double oldM2 = value::bitcastTo<double>(oldM2Val); - - auto [newM2Tag, newM2Val] = nextArr->getAt(AggStdDevValueElems::kRunningM2); - tassert(7039541, "expected double", newM2Tag == value::TypeTags::NumberDouble); - double newM2 = value::bitcastTo<double>(newM2Val); - - const double delta = newMean - oldMean; - // We've already handled the case where 'newCount' is zero above. This means that 'totalCount' - // must be positive, and prevents us from ever dividing by zero in the subsequent calculation. - int64_t totalCount = oldCount + newCount; - if (delta != 0) { - newMean = ((oldCount * oldMean) + (newCount * newMean)) / totalCount; - newM2 += delta * delta * - (static_cast<double>(oldCount) * static_cast<double>(newCount) / totalCount); - } - newM2 += oldM2; - - setStdDevArray(value::bitcastFrom<int64_t>(totalCount), - value::bitcastFrom<double>(newMean), - value::bitcastFrom<double>(newM2), - accumulator); -} - std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggStdDevFinalizeImpl( value::Value fieldValue, bool isSamp) { auto arr = value::getArrayView(fieldValue); @@ -1027,7 +915,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericLn(value::TypeT if (!operand.isGreater(Decimal128::kNormalizedZero) && !operand.isNaN()) { return {false, value::TypeTags::Nothing, 0}; } - auto operandLn = operand.naturalLogarithm(); + auto operandLn = operand.logarithm(); auto [tag, value] = value::makeCopyDecimal(operandLn); return {true, tag, value}; diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp index 812de5232af..7eb8eb7149e 100644 --- a/src/mongo/db/exec/sbe/vm/vm.cpp +++ b/src/mongo/db/exec/sbe/vm/vm.cpp @@ -1072,40 +1072,35 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggSum(value::TypeTags return genericAdd(accTag, accValue, fieldTag, fieldValue); } -template <bool merging> std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAggDoubleDoubleSum( ArityType arity) { + auto [_, fieldTag, fieldValue] = getFromStack(1); // Move the incoming accumulator state from the stack. Given that we are now the owner of the // state we are free to do any in-place update as we see fit. auto [accTag, accValue] = moveOwnedFromStack(0); + value::ValueGuard guard{accTag, accValue}; // Initialize the accumulator. if (accTag == value::TypeTags::Nothing) { std::tie(accTag, accValue) = value::makeNewArray(); - value::ValueGuard newArrGuard{accTag, accValue}; + value::ValueGuard guard{accTag, accValue}; auto arr = value::getArrayView(accValue); arr->reserve(AggSumValueElems::kMaxSizeOfArray); - // The order of the following three elements should match to 'AggSumValueElems'. An absent - // 'kDecimalTotal' element means that we've not seen any decimal value. So, we're not adding - // 'kDecimalTotal' element yet. + // The order of the following three elements should match to 'AggSumValueElems'. arr->push_back(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(0)); arr->push_back(value::TypeTags::NumberDouble, value::bitcastFrom<double>(0.0)); arr->push_back(value::TypeTags::NumberDouble, value::bitcastFrom<double>(0.0)); - newArrGuard.reset(); + // The absent 'kDecimalTotal' element means that we've not seen any decimal value. So, we're + // not adding 'kDecimalTotal' element yet. + aggDoubleDoubleSumImpl(arr, fieldTag, fieldValue); + guard.reset(); + return {true, accTag, accValue}; } - - value::ValueGuard guard{accTag, accValue}; tassert(5755317, "The result slot must be Array-typed", accTag == value::TypeTags::Array); - auto accumulator = value::getArrayView(accValue); - - if constexpr (merging) { - aggMergeDoubleDoubleSumsImpl(accumulator, fieldTag, fieldValue); - } else { - aggDoubleDoubleSumImpl(accumulator, fieldTag, fieldValue); - } + aggDoubleDoubleSumImpl(value::getArrayView(accValue), fieldTag, fieldValue); guard.reset(); return {true, accTag, accValue}; } @@ -1240,37 +1235,31 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinDoubleDoublePar return {true, tag, val}; } -template <bool merging> std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAggStdDev(ArityType arity) { auto [_, fieldTag, fieldValue] = getFromStack(1); // Move the incoming accumulator state from the stack. Given that we are now the owner of the // state we are free to do any in-place update as we see fit. auto [accTag, accValue] = moveOwnedFromStack(0); + value::ValueGuard guard{accTag, accValue}; // Initialize the accumulator. if (accTag == value::TypeTags::Nothing) { - std::tie(accTag, accValue) = value::makeNewArray(); - value::ValueGuard newArrGuard{accTag, accValue}; - auto arr = value::getArrayView(accValue); + auto [newAccTag, newAccValue] = value::makeNewArray(); + value::ValueGuard newGuard{newAccTag, newAccValue}; + auto arr = value::getArrayView(newAccValue); arr->reserve(AggStdDevValueElems::kSizeOfArray); // The order of the following three elements should match to 'AggStdDevValueElems'. arr->push_back(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(0)); arr->push_back(value::TypeTags::NumberDouble, value::bitcastFrom<double>(0.0)); arr->push_back(value::TypeTags::NumberDouble, value::bitcastFrom<double>(0.0)); - newArrGuard.reset(); + aggStdDevImpl(arr, fieldTag, fieldValue); + newGuard.reset(); + return {true, newAccTag, newAccValue}; } - - value::ValueGuard guard{accTag, accValue}; tassert(5755210, "The result slot must be Array-typed", accTag == value::TypeTags::Array); - auto accumulator = value::getArrayView(accValue); - - if constexpr (merging) { - aggMergeStdDevsImpl(accumulator, fieldTag, fieldValue); - } else { - aggStdDevImpl(accumulator, fieldTag, fieldValue); - } + aggStdDevImpl(value::getArrayView(accValue), fieldTag, fieldValue); guard.reset(); return {true, accTag, accValue}; } @@ -3030,120 +3019,6 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinConcat(ArityTyp return {true, strTag, strValue}; } -std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinConcatArrays(ArityType arity) { - auto [resTag, resVal] = value::makeNewArray(); - value::ValueGuard resGuard{resTag, resVal}; - auto resView = value::getArrayView(resVal); - - for (ArityType idx = 0; idx < arity; ++idx) { - auto [_, tag, val] = getFromStack(idx); - if (!value::isArray(tag)) { - return {false, value::TypeTags::Nothing, 0}; - } - - for (auto ae = value::ArrayEnumerator{tag, val}; !ae.atEnd(); ae.advance()) { - auto [elTag, elVal] = ae.getViewOfValue(); - auto [copyTag, copyVal] = value::copyValue(elTag, elVal); - resView->push_back(copyTag, copyVal); - } - } - - resGuard.reset(); - - return {true, resTag, resVal}; -} - -std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAggConcatArraysCapped( - ArityType arity) { - auto [ownArr, tagArr, valArr] = getFromStack(0); - auto [tagNewElem, valNewElem] = moveOwnedFromStack(1); - value::ValueGuard guardNewElem{tagNewElem, valNewElem}; - auto [_, tagSizeCap, valSizeCap] = getFromStack(2); - - tassert(7039508, - "'cap' parameter must be a 32-bit int", - tagSizeCap == value::TypeTags::NumberInt32); - const int32_t sizeCap = value::bitcastTo<int32_t>(valSizeCap); - - // We expect the new value we are adding to the accumulator to be a two-element array where - // the first element is the array to concatenate and the second value is the corresponding size. - tassert(7039512, "expected value of type 'Array'", tagNewElem == value::TypeTags::Array); - auto newArr = value::getArrayView(valNewElem); - tassert(7039527, - "array had unexpected size", - newArr->size() == static_cast<size_t>(AggArrayWithSize::kLast)); - - // Create a new array to hold size and added elements, if is it does not exist yet. - if (tagArr == value::TypeTags::Nothing) { - ownArr = true; - std::tie(tagArr, valArr) = value::makeNewArray(); - auto arr = value::getArrayView(valArr); - - auto [tagAccArr, valAccArr] = value::makeNewArray(); - - // The order is important! The accumulated array should be at index - // AggArrayWithSize::kValues, and the size should be at index - // AggArrayWithSize::kSizeOfValues. - arr->push_back(tagAccArr, valAccArr); - arr->push_back(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(0)); - } else { - // Take ownership of the accumulator. - topStack(false, value::TypeTags::Nothing, 0); - } - - tassert(7039513, "expected array to be owned", ownArr); - value::ValueGuard accumulatorGuard{tagArr, valArr}; - tassert(7039514, "expected accumulator to have type 'Array'", tagArr == value::TypeTags::Array); - auto arr = value::getArrayView(valArr); - tassert(7039515, - "accumulator was array of unexpected size", - arr->size() == static_cast<size_t>(AggArrayWithSize::kLast)); - - // Check that the accumulated size after concatentation won't exceed the limit. - { - auto [tagAccSize, valAccSize] = - arr->getAt(static_cast<size_t>(AggArrayWithSize::kSizeOfValues)); - auto [tagNewSize, valNewSize] = - newArr->getAt(static_cast<size_t>(AggArrayWithSize::kSizeOfValues)); - tassert(7039516, "expected 64-bit int", tagAccSize == value::TypeTags::NumberInt64); - tassert(7039517, "expected 64-bit int", tagNewSize == value::TypeTags::NumberInt64); - const int64_t currentSize = value::bitcastTo<int64_t>(valAccSize); - const int64_t newSize = value::bitcastTo<int64_t>(valNewSize); - const int64_t totalSize = currentSize + newSize; - - if (totalSize >= static_cast<int64_t>(sizeCap)) { - uasserted(ErrorCodes::ExceededMemoryLimit, - str::stream() << "Used too much memory for a single array. Memory limit: " - << sizeCap << ". Concatentating array of " << arr->size() - << " elements and " << currentSize << " bytes with array of " - << newArr->size() << " elements and " << newSize << " bytes."); - } - - // We are still under the size limit. Set the new total size in the accumulator. - arr->setAt(static_cast<size_t>(AggArrayWithSize::kSizeOfValues), - value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(totalSize)); - } - - auto [tagAccArr, valAccArr] = arr->getAt(static_cast<size_t>(AggArrayWithSize::kValues)); - tassert(7039518, "expected value of type 'Array'", tagAccArr == value::TypeTags::Array); - auto accArr = value::getArrayView(valAccArr); - - auto [tagNewArray, valNewArray] = newArr->getAt(static_cast<size_t>(AggArrayWithSize::kValues)); - tassert(7039519, "expected value of type 'Array'", tagNewArray == value::TypeTags::Array); - - for (auto i = value::ArrayEnumerator{tagNewArray, valNewArray}; !i.atEnd(); i.advance()) { - auto [elTag, elVal] = i.getViewOfValue(); - // TODO SERVER-71952: Since 'valNewArray' is owned here, in the future we could avoid this - // copy by moving the element out of the array. - auto [copyTag, copyVal] = value::copyValue(elTag, elVal); - accArr->push_back(copyTag, copyVal); - } - - accumulatorGuard.reset(); - return {ownArr, tagArr, valArr}; -} - std::pair<value::TypeTags, value::Value> ByteCode::genericIsMember(value::TypeTags lhsTag, value::Value lhsVal, value::TypeTags rhsTag, @@ -3521,166 +3396,6 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinSetUnion(ArityT return setUnion(argTags, argVals); } -std::tuple<bool, value::TypeTags, value::Value> ByteCode::aggSetUnionCappedImpl( - value::TypeTags tagNewElem, - value::Value valNewElem, - int32_t sizeCap, - CollatorInterface* collator) { - value::ValueGuard guardNewElem{tagNewElem, valNewElem}; - auto [ownAcc, tagAcc, valAcc] = getFromStack(0); - - // We expect the new value we are adding to the accumulator to be a two-element array where - // the first element is the new set of values and the second value is the corresponding size. - tassert(7039526, "expected value of type 'Array'", tagNewElem == value::TypeTags::Array); - auto newArr = value::getArrayView(valNewElem); - tassert(7039528, - "array had unexpected size", - newArr->size() == static_cast<size_t>(AggArrayWithSize::kLast)); - - // Create a new array is it does not exist yet. - if (tagAcc == value::TypeTags::Nothing) { - ownAcc = true; - std::tie(tagAcc, valAcc) = value::makeNewArray(); - auto accArray = value::getArrayView(valAcc); - - auto [tagAccSet, valAccSet] = value::makeNewArraySet(collator); - - // The order is important! The accumulated array should be at index - // AggArrayWithSize::kValues, and the size should be at index - // AggArrayWithSize::kSizeOfValues. - accArray->push_back(tagAccSet, valAccSet); - accArray->push_back(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(0)); - } else { - // Take ownership of the accumulator. - topStack(false, value::TypeTags::Nothing, 0); - } - - tassert(7039520, "expected accumulator value to be owned", ownAcc); - value::ValueGuard guardArr{tagAcc, valAcc}; - - tassert( - 7039521, "expected accumulator to be of type 'Array'", tagAcc == value::TypeTags::Array); - auto accArray = value::getArrayView(valAcc); - tassert(7039522, - "array had unexpected size", - accArray->size() == static_cast<size_t>(AggArrayWithSize::kLast)); - - auto [tagAccArrSet, valAccArrSet] = - accArray->getAt(static_cast<size_t>(AggArrayWithSize::kValues)); - tassert( - 7039523, "expected value of type 'ArraySet'", tagAccArrSet == value::TypeTags::ArraySet); - auto accArrSet = value::getArraySetView(valAccArrSet); - - // Extract the current size of the accumulator. As we add elements to the set, we will increment - // the current size accordingly and throw an exception if we ever exceed the size limit. We - // cannot simply sum the two sizes, since the two sets could have a substantial intersection. - auto [tagAccSize, valAccSize] = - accArray->getAt(static_cast<size_t>(AggArrayWithSize::kSizeOfValues)); - tassert(7039524, "expected 64-bit int", tagAccSize == value::TypeTags::NumberInt64); - int64_t currentSize = value::bitcastTo<int64_t>(valAccSize); - - auto [tagNewValSet, valNewValSet] = - newArr->getAt(static_cast<size_t>(AggArrayWithSize::kValues)); - tassert( - 7039525, "expected value of type 'ArraySet'", tagNewValSet == value::TypeTags::ArraySet); - - for (auto i = value::ArrayEnumerator{tagNewValSet, valNewValSet}; !i.atEnd(); i.advance()) { - auto [elTag, elVal] = i.getViewOfValue(); - int elemSize = value::getApproximateSize(elTag, elVal); - // TODO SERVER-71952: Since 'valNewValSet' is owned here, in the future we could avoid this - // copy by moving the element out of the array. - auto [copyTag, copyVal] = value::copyValue(elTag, elVal); - bool inserted = accArrSet->push_back(copyTag, copyVal); - - if (inserted) { - currentSize += elemSize; - if (currentSize >= static_cast<int64_t>(sizeCap)) { - uasserted(ErrorCodes::ExceededMemoryLimit, - str::stream() << "Used too much memory for a single array. Memory limit: " - << sizeCap << ". Current set has " << accArrSet->size() - << " elements and is " << currentSize << " bytes."); - } - } - } - - // Update the accumulator with the new total size. - accArray->setAt(static_cast<size_t>(AggArrayWithSize::kSizeOfValues), - value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(currentSize)); - - guardArr.reset(); - return {ownAcc, tagAcc, valAcc}; -} - -std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAggSetUnion(ArityType arity) { - auto [ownAcc, tagAcc, valAcc] = getFromStack(0); - - if (tagAcc == value::TypeTags::Nothing) { - // Initialize the accumulator. - ownAcc = true; - std::tie(tagAcc, valAcc) = value::makeNewArraySet(); - } else { - // Take ownership of the accumulator. - topStack(false, value::TypeTags::Nothing, 0); - } - - tassert(7039552, "accumulator must be owned", ownAcc); - value::ValueGuard guardAcc{tagAcc, valAcc}; - tassert(7039553, "accumulator must be of type ArraySet", tagAcc == value::TypeTags::ArraySet); - auto acc = value::getArraySetView(valAcc); - - auto [tagNewSet, valNewSet] = moveOwnedFromStack(1); - value::ValueGuard guardNewSet{tagNewSet, valNewSet}; - if (!value::isArray(tagNewSet)) { - return {false, value::TypeTags::Nothing, 0}; - } - - auto i = value::ArrayEnumerator{tagNewSet, valNewSet}; - while (!i.atEnd()) { - auto [elTag, elVal] = i.getViewOfValue(); - auto [copyTag, copyVal] = value::copyValue(elTag, elVal); - acc->push_back(copyTag, copyVal); - i.advance(); - } - - guardAcc.reset(); - return {ownAcc, tagAcc, valAcc}; -} - -std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAggSetUnionCapped( - ArityType arity) { - auto [tagNewElem, valNewElem] = moveOwnedFromStack(1); - value::ValueGuard guardNewElem{tagNewElem, valNewElem}; - - auto [_, tagSizeCap, valSizeCap] = getFromStack(2); - tassert(7039509, - "'cap' parameter must be a 32-bit int", - tagSizeCap == value::TypeTags::NumberInt32); - const size_t sizeCap = value::bitcastTo<int32_t>(valSizeCap); - - guardNewElem.reset(); - return aggSetUnionCappedImpl(tagNewElem, valNewElem, sizeCap, nullptr /*collator*/); -} - -std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinAggCollSetUnionCapped( - ArityType arity) { - auto [_1, tagColl, valColl] = getFromStack(1); - auto [tagNewElem, valNewElem] = moveOwnedFromStack(2); - value::ValueGuard guardNewElem{tagNewElem, valNewElem}; - auto [_2, tagSizeCap, valSizeCap] = getFromStack(3); - - tassert(7039510, "expected value of type 'collator'", tagColl == value::TypeTags::collator); - tassert(7039511, - "'cap' parameter must be a 32-bit int", - tagSizeCap == value::TypeTags::NumberInt32); - - guardNewElem.reset(); - return aggSetUnionCappedImpl(tagNewElem, - valNewElem, - value::bitcastTo<int32_t>(valSizeCap), - value::getCollatorView(valColl)); -} - std::tuple<bool, value::TypeTags, value::Value> ByteCode::builtinCollSetIntersection( ArityType arity) { invariant(arity >= 1); @@ -4667,7 +4382,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builti case Builtin::doubleDoubleSum: return builtinDoubleDoubleSum(arity); case Builtin::aggDoubleDoubleSum: - return builtinAggDoubleDoubleSum<false /*merging*/>(arity); + return builtinAggDoubleDoubleSum(arity); case Builtin::doubleDoubleSumFinalize: return builtinDoubleDoubleSumFinalize<>(arity); case Builtin::doubleDoubleMergeSumFinalize: @@ -4676,12 +4391,8 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builti return builtinDoubleDoubleSumFinalize<true /*keepIntegerPrecision*/>(arity); case Builtin::doubleDoublePartialSumFinalize: return builtinDoubleDoublePartialSumFinalize(arity); - case Builtin::aggMergeDoubleDoubleSums: - return builtinAggDoubleDoubleSum<true /*merging*/>(arity); case Builtin::aggStdDev: - return builtinAggStdDev<false /*merging*/>(arity); - case Builtin::aggMergeStdDevs: - return builtinAggStdDev<true /*merging*/>(arity); + return builtinAggStdDev(arity); case Builtin::stdDevPopFinalize: return builtinStdDevPopFinalize(arity); case Builtin::stdDevSampFinalize: @@ -4734,16 +4445,6 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builti return builtinRound(arity); case Builtin::concat: return builtinConcat(arity); - case Builtin::concatArrays: - return builtinConcatArrays(arity); - case Builtin::aggConcatArraysCapped: - return builtinAggConcatArraysCapped(arity); - case Builtin::aggSetUnion: - return builtinAggSetUnion(arity); - case Builtin::aggSetUnionCapped: - return builtinAggSetUnionCapped(arity); - case Builtin::aggCollSetUnionCapped: - return builtinAggCollSetUnionCapped(arity); case Builtin::isMember: return builtinIsMember(arity); case Builtin::collIsMember: diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h index e5148e8d4b2..255a1a497c2 100644 --- a/src/mongo/db/exec/sbe/vm/vm.h +++ b/src/mongo/db/exec/sbe/vm/vm.h @@ -511,33 +511,12 @@ enum class Builtin : uint8_t { collAddToSet, // agg function to append to a set (with collation) collAddToSetCapped, // agg function to append to a set (with collation), fails when the set // reaches specified size - - // Special double summation. - doubleDoubleSum, - // A variant of the standard sum aggregate function which maintains a DoubleDouble as the - // accumulator's underlying state. + doubleDoubleSum, // special double summation aggDoubleDoubleSum, - // Converts a DoubleDouble sum into a single numeric scalar for use once the summation is - // complete. doubleDoubleSumFinalize, - // A form of doubleDoubleSum finalization only necessary for sharding support when the cluster - // is not yet fully upgraded to FCV 6.0. doubleDoubleMergeSumFinalize, - // Converts a partial sum into a format suitable for serialization over the wire to the merging - // node. The merging node expects the internal state of the DoubleDouble summation to be - // serialized in a particular format. doubleDoublePartialSumFinalize, - // An agg function which can be used to sum a sequence of DoubleDouble inputs, producing the - // resulting total as a DoubleDouble. - aggMergeDoubleDoubleSums, - - // Implements Welford's online algorithm for computing sample or population standard deviation - // in a single pass. aggStdDev, - // Combines standard deviations that have been partially computed on a subset of the data - // using Welford's online algorithm. - aggMergeStdDevs, - stdDevPopFinalize, stdDevSampFinalize, bitTestZero, // test bitwise mask & value is zero @@ -548,18 +527,6 @@ enum class Builtin : uint8_t { toLower, coerceToString, concat, - concatArrays, - - // Agg function to concatenate arrays, failing when the accumulator reaches a specified size. - aggConcatArraysCapped, - - // Agg functions to compute the set union of two arrays, failing when the accumulator reaches a - // specified size. - aggSetUnionCapped, - aggCollSetUnionCapped, - // Agg function for a simple set union (with no size cap or collation). - aggSetUnion, - acos, acosh, asin, @@ -1013,19 +980,11 @@ private: value::TypeTags fieldTag, value::Value fieldValue); - void aggDoubleDoubleSumImpl(value::Array* accumulator, - value::TypeTags rhsTag, - value::Value rhsValue); - void aggMergeDoubleDoubleSumsImpl(value::Array* accumulator, - value::TypeTags rhsTag, - value::Value rhsValue); + void aggDoubleDoubleSumImpl(value::Array* arr, value::TypeTags rhsTag, value::Value rhsValue); // This is an implementation of the following algorithm: // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm - void aggStdDevImpl(value::Array* accumulator, value::TypeTags rhsTag, value::Value rhsValue); - void aggMergeStdDevsImpl(value::Array* accumulator, - value::TypeTags rhsTag, - value::Value rhsValue); + void aggStdDevImpl(value::Array* arr, value::TypeTags rhsTag, value::Value rhsValue); std::tuple<bool, value::TypeTags, value::Value> aggStdDevFinalizeImpl(value::Value fieldValue, bool isSamp); @@ -1144,24 +1103,14 @@ private: CollatorInterface* collator); std::tuple<bool, value::TypeTags, value::Value> builtinAddToSetCapped(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinCollAddToSetCapped(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> builtinDoubleDoubleSum(ArityType arity); - // The template parameter is false for a regular DoubleDouble summation and true if merging - // partially computed DoubleDouble sums. - template <bool merging> std::tuple<bool, value::TypeTags, value::Value> builtinAggDoubleDoubleSum(ArityType arity); - // This is only for compatibility with mongos/sharding and we will revisit this later. template <bool keepIntegerPrecision = false> std::tuple<bool, value::TypeTags, value::Value> builtinDoubleDoubleSumFinalize(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinDoubleDoublePartialSumFinalize( ArityType arity); - - // The template parameter is false for a regular std dev and true if merging partially computed - // standard devations. - template <bool merging> std::tuple<bool, value::TypeTags, value::Value> builtinAggStdDev(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> builtinStdDevPopFinalize(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinStdDevSampFinalize(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinBitTestZero(ArityType arity); @@ -1188,16 +1137,6 @@ private: std::tuple<bool, value::TypeTags, value::Value> builtinTanh(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinRound(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinConcat(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> builtinConcatArrays(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> builtinAggConcatArraysCapped(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> builtinAggSetUnion(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> builtinAggSetUnionCapped(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> builtinAggCollSetUnionCapped(ArityType arity); - std::tuple<bool, value::TypeTags, value::Value> aggSetUnionCappedImpl( - value::TypeTags tagNewElem, - value::Value valNewElem, - int32_t sizeCap, - CollatorInterface* collator); std::tuple<bool, value::TypeTags, value::Value> builtinIsMember(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinCollIsMember(ArityType arity); std::tuple<bool, value::TypeTags, value::Value> builtinIndexOfBytes(ArityType arity); |