diff options
Diffstat (limited to 'src/mongo/db/exec/sbe')
33 files changed, 1420 insertions, 475 deletions
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript index b99625b0833..f8246baf556 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/key_generator', + '$BUILD_DIR/mongo/db/index/index_access_method', '$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,8 +59,9 @@ 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', @@ -89,7 +90,6 @@ 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,6 +106,7 @@ 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 7da6a8e2cef..7939d166a78 100644 --- a/src/mongo/db/exec/sbe/abt/abt_lower.cpp +++ b/src/mongo/db/exec/sbe/abt/abt_lower.cpp @@ -586,14 +586,15 @@ std::unique_ptr<sbe::PlanStage> SBENodeLowering::walk(const GroupByNode& n, auto& names = binderAgg->names(); auto& exprs = refsAgg->nodes(); - sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> aggs; + sbe::SlotExprPairVector aggs; + aggs.reserve(exprs.size()); 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.emplace(slot, std::move(expr)); + aggs.push_back({slot, std::move(expr)}); } // TODO: use collator slot. @@ -609,6 +610,11 @@ 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); } @@ -1005,6 +1011,7 @@ 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 72d9820d760..bc5121407f9 100644 --- a/src/mongo/db/exec/sbe/expression_test_base.h +++ b/src/mongo/db/exec/sbe/expression_test_base.h @@ -69,6 +69,24 @@ 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 3a2b20a4657..65421e8b373 100644 --- a/src/mongo/db/exec/sbe/expressions/expression.cpp +++ b/src/mongo/db/exec/sbe/expressions/expression.cpp @@ -429,6 +429,8 @@ 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", @@ -436,6 +438,8 @@ 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", @@ -467,7 +471,10 @@ 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{[](size_t n) { return n > 0; }, vm::Builtin::concat, 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}}, {"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", @@ -486,6 +493,11 @@ 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 2b9032f257b..d5a2c0fdf0e 100644 --- a/src/mongo/db/exec/sbe/expressions/expression.h +++ b/src/mongo/db/exec/sbe/expressions/expression.h @@ -348,6 +348,8 @@ 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)...); @@ -364,20 +366,30 @@ inline auto makeEs(Ts&&... pack) { namespace detail { // base case -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)); +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)}); + } } // recursive case -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)...); +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)...); } } // namespace detail @@ -386,7 +398,7 @@ auto makeEM(Ts&&... pack) { value::SlotMap<std::unique_ptr<EExpression>> result; if constexpr (sizeof...(pack) > 0) { result.reserve(sizeof...(Ts) / 2); - detail::makeEM_unwind(result, std::forward<Ts>(pack)...); + detail::makeSlotExprPairHelper(result, std::forward<Ts>(pack)...); } return result; } @@ -399,6 +411,16 @@ 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 30eb61a7b2d..279f7a287b4 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,6 +28,7 @@ */ #include "mongo/db/exec/sbe/expression_test_base.h" +#include "mongo/db/query/sbe_stage_builder_helpers.h" namespace mongo::sbe { @@ -93,6 +94,35 @@ 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 87fc8dbd1f2..89244fd4589 100644 --- a/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp +++ b/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp @@ -69,18 +69,23 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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)); @@ -166,20 +171,21 @@ TEST_F(HashAggStageTest, HashAggMinMaxTest) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(), - 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))), + 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))), makeSV(), true, boost::none, false /* allowDiskUse */, + makeSlotExprPairVec() /* mergingExprs */, kEmptyPlanNodeId); auto outSlot = generateSlotId(); @@ -230,13 +236,15 @@ TEST_F(HashAggStageTest, HashAggAddToSetTest) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(), - makeEM(hashAggSlot, - stage_builder::makeFunction( - "collAddToSet", std::move(collExpr), makeE<EVariable>(scanSlot))), + makeSlotExprPairVec(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)); @@ -327,14 +335,16 @@ TEST_F(HashAggStageTest, HashAggSeekKeysTest) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeEM(countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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)); @@ -387,17 +397,21 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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. @@ -420,6 +434,7 @@ 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(); @@ -428,7 +443,6 @@ 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); @@ -446,17 +460,21 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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. @@ -479,7 +497,14 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpill) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_TRUE(stats->usedDisk); - ASSERT_EQ(results.size() - expectedRowsToFitInMemory, stats->spilledRecords); + // 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); stage->close(); } @@ -513,17 +538,21 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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. @@ -546,6 +575,7 @@ 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(); @@ -554,7 +584,6 @@ 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); @@ -572,17 +601,21 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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. @@ -605,7 +638,14 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpillDouble) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_TRUE(stats->usedDisk); - ASSERT_EQ(results.size() - expectedRowsToFitInMemory, stats->spilledRecords); + // 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); stage->close(); } @@ -627,17 +667,21 @@ 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(), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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. @@ -658,6 +702,7 @@ 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(); @@ -666,7 +711,6 @@ 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); @@ -685,19 +729,27 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1))), - sumsSlot, - stage_builder::makeFunction("sum", makeE<EVariable>(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))), 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. @@ -727,7 +779,10 @@ 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(results.size() - expectedRowsToFitInMemory, stats->spilledRecords); + 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); stage->close(); } @@ -752,19 +807,27 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1))), - sumsSlot, - stage_builder::makeFunction("sum", makeE<EVariable>(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))), 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. @@ -794,7 +857,9 @@ TEST_F(HashAggStageTest, HashAggMultipleAccSpillAllToDisk) { // Check that the spilling behavior matches the expected. auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats()); ASSERT_TRUE(stats->usedDisk); - ASSERT_EQ(results.size(), stats->spilledRecords); + // We expect each incoming value to result in a spill of a single record. + ASSERT_EQ(stats->numSpills, 9); + ASSERT_EQ(stats->spilledRecords, 9); stage->close(); } @@ -831,14 +896,18 @@ 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), - makeEM(sumsSlot, stage_builder::makeFunction("sum", makeE<EVariable>(scanSlot))), + makeSlotExprPairVec(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. @@ -872,17 +941,21 @@ 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), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), + makeSlotExprPairVec( + 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 191f0ffe562..f62b8a6229c 100644 --- a/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp +++ b/src/mongo/db/exec/sbe/sbe_plan_size_test.cpp @@ -139,11 +139,12 @@ TEST_F(PlanSizeTest, Filter) { TEST_F(PlanSizeTest, HashAgg) { auto stage = makeS<HashAggStage>(mockS(), mockSV(), - makeEM(generateSlotId(), mockE()), + makeSlotExprPairVec(generateSlotId(), mockE()), makeSV(), true, generateSlotId(), false, + makeSlotExprPairVec(), kEmptyPlanNodeId); assertPlanSize(*stage); } @@ -168,6 +169,7 @@ TEST_F(PlanSizeTest, IndexScan) { generateSlotId(), generateSlotId(), generateSlotId(), + generateSlotId(), IndexKeysInclusionSet(1), mockSV(), generateSlotId(), 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 54fadbe1f15..57d3ff4e05f 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,14 +141,16 @@ TEST_F(TrialRunTrackerTest, TrialEndsDuringOpenPhaseOfBlockingStage) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), - makeSV(), // Seek slot + makeSlotExprPairVec( + 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}); @@ -210,14 +212,16 @@ TEST_F(TrialRunTrackerTest, OnlyDeepestNestedBlockingStageHasTrialRunTracker) { auto hashAggStage = makeS<HashAggStage>( std::move(unionStage), makeSV(unionSlot), - makeEM(countsSlot, - stage_builder::makeFunction( - "sum", - makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))), - makeSV(), // Seek slot + makeSlotExprPairVec( + 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); @@ -277,14 +281,16 @@ TEST_F(TrialRunTrackerTest, SiblingBlockingStagesBothGetTrialRunTracker) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeEM(countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), - makeSV(), // Seek slot + makeSlotExprPairVec( + 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)); @@ -407,14 +413,16 @@ TEST_F(TrialRunTrackerTest, DisablingTrackingForAChildStagePreventsEarlyExit) { auto hashAggStage = makeS<HashAggStage>( std::move(scanStage), makeSV(scanSlot), - makeEM(countsSlot, - stage_builder::makeFunction("sum", - makeE<EConstant>(value::TypeTags::NumberInt64, - value::bitcastFrom<int64_t>(1)))), - makeSV(), // Seek slot + makeSlotExprPairVec( + 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 fb6684eea52..bbb92328331 100644 --- a/src/mongo/db/exec/sbe/size_estimator.h +++ b/src/mongo/db/exec/sbe/size_estimator.h @@ -92,6 +92,11 @@ 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 bec12b12ee2..1e8809f0858 100644 --- a/src/mongo/db/exec/sbe/stages/branch.cpp +++ b/src/mongo/db/exec/sbe/stages/branch.cpp @@ -70,31 +70,29 @@ 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) { - std::vector<value::SlotAccessor*> accessors; - accessors.reserve(2); + auto slot = _outputVals[idx]; + auto [_, inserted] = dupCheck.insert(slot); + uassert(4822831, str::stream() << "duplicate field: " << slot, inserted); + } - { - auto slot = _inputThenVals[idx]; - auto [it, inserted] = dupCheck.insert(slot); - uassert(4822829, str::stream() << "duplicate field: " << slot, inserted); + for (size_t idx = 0; idx < _outputVals.size(); ++idx) { + auto thenSlot = _inputThenVals[idx]; + auto elseSlot = _inputElseVals[idx]; - 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); + // 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[1]->getAccessor(ctx, slot)); - } - { - auto slot = _outputVals[idx]; - auto [it, inserted] = dupCheck.insert(slot); - uassert(4822831, 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)); - _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 4116f2980ff..7471763a44a 100644 --- a/src/mongo/db/exec/sbe/stages/collection_helpers.h +++ b/src/mongo/db/exec/sbe/stages/collection_helpers.h @@ -40,6 +40,7 @@ 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 14f40afeaa9..bdbcda14805 100644 --- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp +++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp @@ -27,47 +27,67 @@ * it in the license file. */ -#include "mongo/platform/basic.h" +#include "mongo/db/exec/sbe/stages/hash_agg.h" #include "mongo/db/concurrency/d_concurrency.h" -#include "mongo/db/concurrency/write_conflict_exception.h" -#include "mongo/db/exec/sbe/stages/hash_agg.h" +#include "mongo/db/exec/sbe/size_estimator.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, - value::SlotMap<std::unique_ptr<EExpression>> aggs, + SlotExprPairVector aggs, value::SlotVector seekKeysSlots, bool optimizedClose, boost::optional<value::SlotId> collatorSlot, bool allowDiskUse, - PlanNodeId planNodeId) + SlotExprPairVector mergingExprs, + PlanNodeId planNodeId, + bool forceIncreasedSpilling) : PlanStage("group"_sd, planNodeId), _gbs(std::move(gbs)), _aggs(std::move(aggs)), _collatorSlot(collatorSlot), _allowDiskUse(allowDiskUse), _seekKeysSlots(std::move(seekKeysSlots)), - _optimizedClose(optimizedClose) { + _optimizedClose(optimizedClose), + _mergingExprs(std::move(mergingExprs)), + _forceIncreasedSpilling(forceIncreasedSpilling) { _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 { - value::SlotMap<std::unique_ptr<EExpression>> aggs; + SlotExprPairVector aggs; + aggs.reserve(_aggs.size()); for (auto& [k, v] : _aggs) { - aggs.emplace(k, v->clone()); + aggs.push_back({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), @@ -75,7 +95,9 @@ std::unique_ptr<PlanStage> HashAggStage::clone() const { _optimizedClose, _collatorSlot, _allowDiskUse, - _commonStats.nodeId); + std::move(mergingExprsClone), + _commonStats.nodeId, + _forceIncreasedSpilling); } void HashAggStage::doSaveState(bool relinquishCursor) { @@ -120,34 +142,36 @@ 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) { - auto [it, inserted] = dupCheck.emplace(slot); - uassert(4822827, str::stream() << "duplicate field: " << slot, inserted); + throwIfDupSlot(slot); _inKeyAccessors.emplace_back(_children[0]->getAccessor(ctx, slot)); - // 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. + // Construct accessors for obtaining the key values from either the hash table '_ht' or the + // '_recordStore'. _outHashKeyAccessors.emplace_back(std::make_unique<HashKeyAccessor>(_htIt, counter)); _outRecordStoreKeyAccessors.emplace_back( - std::make_unique<value::MaterializedSingleRowAccessor>(_aggKeyRecordStore, counter)); - - counter++; + std::make_unique<value::MaterializedSingleRowAccessor>(_outKeyRowRecordStore, 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. 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'. + // 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. _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 @@ -158,26 +182,18 @@ void HashAggStage::prepare(CompileCtx& ctx) { counter = 0; for (auto& [slot, expr] : _aggs) { - 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. + 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'. _outRecordStoreAggAccessors.emplace_back( - std::make_unique<value::MaterializedSingleRowAccessor>(_aggValueRecordStore, counter)); + std::make_unique<value::MaterializedSingleRowAccessor>(_outAggRowRecordStore, counter)); _outHashAggAccessors.emplace_back(std::make_unique<HashAggAccessor>(_htIt, counter)); - 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. + + // 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. _outAggAccessors.emplace_back( std::make_unique<value::SwitchAccessor>(std::vector<value::SlotAccessor*>{ _outHashAggAccessors.back().get(), _outRecordStoreAggAccessors.back().get()})); @@ -187,10 +203,32 @@ 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; } @@ -200,6 +238,15 @@ 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); } @@ -225,89 +272,91 @@ void HashAggStage::spillRowToDisk(const value::MaterializedRow& key, const value::MaterializedRow& val) { KeyString::Builder kb{KeyString::Version::kLatestVersion}; key.serializeIntoKeyString(kb); + // 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 typeBits = kb.getTypeBits(); - 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); - spillValueToDisk(rid, val, typeBits, false /*update*/); + upsertToRecordStore(_opCtx, _recordStore->rs(), rid, val, typeBits, false /*update*/); + _specificStats.spilledRecords++; } -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++; +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(); } - _specificStats.lastSpilledRecordSize = nBytes; + + for (auto&& it : *_ht) { + spillRowToDisk(it.first, it.second); + } + + 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; } // 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 (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. +// spilling. void HashAggStage::checkMemoryUsageAndSpillIfNecessary(MemoryCheckData& mcd) { - // 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; - } + invariant(!_ht->empty()); // If the group-by key is empty we will only ever aggregate into a single row so no sense in - // spilling since we will just be moving a single row back and forth from disk to main memory. + // spilling. if (_inKeyAccessors.size() == 0) { return; } mcd.memoryCheckpointCounter++; - 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; + 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. 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 - : (estimatedGainPerChildAdvance < -0.1) ? mcd.atMostCheckFrequency - : mcd.nextMemoryCheckpoint * 2; + : mcd.nextMemoryCheckpoint * 2; + mcd.nextMemoryCheckpoint = std::min<long>(mcd.memoryCheckFrequency, std::max<long>(mcd.atMostCheckFrequency, nextCheckpointCandidate)); @@ -333,17 +382,28 @@ 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); - const value::MaterializedRowEq equator(collatorView); - _ht.emplace(0, hasher, equator); + _keyEq = value::MaterializedRowEq(collatorView); + _ht.emplace(0, hasher, _keyEq); } else { _ht.emplace(); } _seekKeys.resize(_seekKeysAccessors.size()); - // A default value for spilling a key to the record store. - value::MaterializedRow defaultVal{_outAggAccessors.size()}; - bool updateAggStateHt = false; + // Reset state since this stage may have been previously opened. + for (auto&& accessor : _outKeyAccessors) { + accessor->setIndex(0); + } + for (auto&& accessor : _outAggAccessors) { + accessor->setIndex(0); + } + _rsCursor.reset(); + _recordStore.reset(); + _outKeyRowRecordStore = {0}; + _outAggRowRecordStore = {0}; + _spilledAggRow = {0}; + _stashedNextRow = {0, 0}; + MemoryCheckData memoryCheckData; while (_children[0]->getNext() == PlanState::ADVANCED) { @@ -355,56 +415,34 @@ void HashAggStage::open(bool reOpen) { key.reset(idx++, false, tag, val); } - - 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. + 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; key.makeOwned(); auto [it, _] = _ht->emplace(std::move(key), value::MaterializedRow{0}); - // Initialize accumulators. it->second.resize(_outAggAccessors.size()); _htIt = it; } - 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 { - // 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); + + // 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); } - // Estimates how much memory is being used and might start spilling. - checkMemoryUsageAndSpillIfNecessary(memoryCheckData); + if (_forceIncreasedSpilling && !newKey) { + // If configured to spill more than usual, we spill after seeing the same key twice. + spill(memoryCheckData); + } 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); + } if (_tracker && _tracker->trackProgress<TrialRunTracker::kNumResults>(1)) { // During trial runs, we want to limit the amount of work done by opening a blocking @@ -421,8 +459,33 @@ 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->rs()->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. @@ -434,20 +497,76 @@ 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, {}); + 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)}; +} - // 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); +PlanState HashAggStage::getNextSpilled() { + 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 = deserializeSpilledRecord(*nextRecord, _outKeyRowRSBuffer); + + _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}; } - _drainingRecordStore = false; + + // 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 = deserializeSpilledRecord(*nextRecord, _stashedKeyBuffer); + 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); + } + } + + return trackPlanState(PlanState::ADVANCED); } PlanState HashAggStage::getNext() { auto optTimer(getOptTimer(_opCtx)); - if (_htIt == _ht->end() && !_drainingRecordStore) { - // First invocation of getNext() after open() when not draining the '_recordStore'. + // 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 (!_seekKeysAccessors.empty()) { _htIt = _ht->find(_seekKeys); } else { @@ -456,53 +575,13 @@ PlanState HashAggStage::getNext() { } else if (!_seekKeysAccessors.empty()) { // Subsequent invocation with seek keys. Return only 1 single row (if any). _htIt = _ht->end(); - } else if (!_drainingRecordStore) { - // Returning the results of the entire hash table first before draining the '_recordStore'. + } else { ++_htIt; } - if (_htIt == _ht->end() && !_recordStore) { - // The hash table has been drained and nothing was spilled to disk. + if (_htIt == _ht->end()) { + // The hash table has been drained (and we never spilled to disk) so we're done. 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); } @@ -522,11 +601,19 @@ 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("spilledBytesApprox", - _specificStats.lastSpilledRecordSize * _specificStats.spilledRecords); + bob.appendNumber("spilledDataStorageSize", _specificStats.spilledDataStorageSize); ret->debugInfo = bob.obj(); } @@ -544,11 +631,12 @@ void HashAggStage::close() { trackClose(); _ht = boost::none; - 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(); @@ -571,7 +659,7 @@ std::vector<DebugPrinter::Block> HashAggStage::debugPrint() const { ret.emplace_back(DebugPrinter::Block("[`")); bool first = true; - value::orderedSlotMapTraverse(_aggs, [&](auto slot, auto&& expr) { + for (auto&& [slot, expr] : _aggs) { if (!first) { ret.emplace_back(DebugPrinter::Block("`,")); } @@ -580,7 +668,7 @@ std::vector<DebugPrinter::Block> HashAggStage::debugPrint() const { ret.emplace_back("="); DebugPrinter::addBlocks(ret, expr->debugPrint()); first = false; - }); + } ret.emplace_back("`]"); if (!_seekKeysSlots.empty()) { @@ -595,6 +683,28 @@ 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"); } @@ -615,6 +725,7 @@ 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 19fbca9d1c7..001b29be887 100644 --- a/src/mongo/db/exec/sbe/stages/hash_agg.h +++ b/src/mongo/db/exec/sbe/stages/hash_agg.h @@ -29,8 +29,6 @@ #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/vm/vm.h" @@ -61,21 +59,38 @@ 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>]? reopen? - * collatorSlot? childStage + * 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 */ class HashAggStage final : public PlanStage { public: HashAggStage(std::unique_ptr<PlanStage> input, value::SlotVector gbs, - value::SlotMap<std::unique_ptr<EExpression>> aggs, + SlotExprPairVector aggs, value::SlotVector seekKeysSlots, bool optimizedClose, boost::optional<value::SlotId> collatorSlot, bool allowDiskUse, - PlanNodeId planNodeId); + SlotExprPairVector mergingExprs, + PlanNodeId planNodeId, + bool forceIncreasedSpilling = false); std::unique_ptr<PlanStage> clone() const final; @@ -108,99 +123,138 @@ private: using HashKeyAccessor = value::MaterializedRowKeyAccessor<TableType::iterator>; using HashAggAccessor = value::MaterializedRowValueAccessor<TableType::iterator>; - 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); + using SpilledRow = std::pair<value::MaterializedRow, value::MaterializedRow>; /** * 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' and might spill some of - * the already accumulated data into it. + * 'checkMemoryUsageAndSpillIfNecessary()' will create '_recordStore' (if it hasn't already been + * created) and spill the contents of the hash table into this record store. */ struct MemoryCheckData { + MemoryCheckData() { + reset(); + } + + void reset() { + memoryCheckFrequency = std::min(atMostCheckFrequency, atLeastMemoryCheckFrequency); + nextMemoryCheckpoint = 0; + memoryCheckpointCounter = 0; + lastEstimatedMemoryUsage = 0; + } + const double checkpointMargin = internalQuerySBEAggMemoryUseCheckMargin.load(); - const long atMostCheckFrequency = internalQuerySBEAggMemoryCheckPerAdvanceAtMost.load(); - const long atLeastMemoryCheckFrequency = + const int64_t atMostCheckFrequency = internalQuerySBEAggMemoryCheckPerAdvanceAtMost.load(); + const int64_t 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. - long memoryCheckFrequency = 1; + int64_t memoryCheckFrequency = 1; // The number of incoming records to process before the next memory checkpoint. - long nextMemoryCheckpoint = 0; + int64_t nextMemoryCheckpoint = 0; // The counter of the incoming records between memory checkpoints. - long memoryCheckpointCounter = 0; + int64_t memoryCheckpointCounter = 0; - long long lastEstimatedMemoryUsage = 0; - - MemoryCheckData() { - memoryCheckFrequency = std::min(atMostCheckFrequency, atLeastMemoryCheckFrequency); - } + int64_t lastEstimatedMemoryUsage = 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); + 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. + */ + SpilledRow deserializeSpilledRecord(const Record& record, BufBuilder& keyBuffer); + + PlanState getNextSpilled(); + + void makeTemporaryRecordStore(); const value::SlotVector _gbs; - const value::SlotMap<std::unique_ptr<EExpression>> _aggs; + const SlotExprPairVector _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; - // Accessors for the key stored in '_ht', a SwitchAccessor is used so we can produce the key - // from either the '_ht' or the '_recordStore'. + // 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. 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; - // 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; + // 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}; std::vector<std::unique_ptr<value::MaterializedSingleRowAccessor>> _outRecordStoreAggAccessors; - - // 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<std::unique_ptr<value::SwitchAccessor>> _outAggAccessors; std::vector<value::SlotAccessor*> _seekKeysAccessors; value::MaterializedRow _seekKeys; - // 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; + // Bytecode which gets executed to aggregate incoming rows into the hash table. 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; @@ -212,10 +266,31 @@ 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<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 0d976f1a875..fbc8ff73058 100644 --- a/src/mongo/db/exec/sbe/stages/hash_lookup.cpp +++ b/src/mongo/db/exec/sbe/stages/hash_lookup.cpp @@ -334,7 +334,7 @@ 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 (newMemUsage <= _memoryUseInBytesBeforeSpill) { + if (!hasSpilledBufToDisk() && newMemUsage <= _memoryUseInBytesBeforeSpill) { _buffer.emplace_back(std::move(value)); _computedTotalMemUsage = newMemUsage; } else { diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp index bfad6d9a2ae..fe88d9c9095 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp +++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp @@ -45,6 +45,7 @@ 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, @@ -58,6 +59,7 @@ IndexScanStage::IndexScanStage(UUID collUuid, _recordSlot(recordSlot), _recordIdSlot(recordIdSlot), _snapshotIdSlot(snapshotIdSlot), + _indexIdSlot(indexIdSlot), _indexKeysToInclude(indexKeysToInclude), _vars(std::move(vars)), _seekKeySlotLow(seekKeySlotLow), @@ -76,6 +78,7 @@ std::unique_ptr<PlanStage> IndexScanStage::clone() const { _recordSlot, _recordIdSlot, _snapshotIdSlot, + _indexIdSlot, _indexKeysToInclude, _vars, _seekKeySlotLow, @@ -128,10 +131,17 @@ 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) { - _snapshotIdAccessor->reset( - value::TypeTags::NumberInt64, - value::bitcastFrom<uint64_t>(_opCtx->recoveryUnit()->getSnapshotId().toNumber())); + _latestSnapshotId = _opCtx->recoveryUnit()->getSnapshotId().toNumber(); } } @@ -148,6 +158,10 @@ 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; } @@ -219,9 +233,7 @@ 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) { - _snapshotIdAccessor->reset( - value::TypeTags::NumberInt64, - value::bitcastFrom<uint64_t>(_opCtx->recoveryUnit()->getSnapshotId().toNumber())); + _latestSnapshotId = _opCtx->recoveryUnit()->getSnapshotId().toNumber(); } } @@ -385,6 +397,12 @@ 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( @@ -423,6 +441,9 @@ 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)); } @@ -471,6 +492,12 @@ 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 ce00ef17128..33fcd352d0c 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.h +++ b/src/mongo/db/exec/sbe/stages/ix_scan.h @@ -51,7 +51,8 @@ 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, and + * - 'snapshotIdSlot': the storage snapshot that this index scan is reading from, + * - 'indexIdSlot': the name of the index being read 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 @@ -63,10 +64,11 @@ namespace mongo::sbe { * * Debug string representation: * - * ixscan recordSlot? recordIdSlot? snapshotIdSlot? [slot_1 = fieldNo_1, ..., slot2 = fieldNo_n] + * ixscan recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? + * [slot_1 = fieldNo_1, ..., slot2 = fieldNo_n] * collectionUuid indexName forward * - * ixseek lowKey highKey recordSlot? recordIdSlot? snapshotIdSlot? + * ixseek lowKey highKey recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? * [slot_1 = fieldNo_1, ..., slot2 = fieldNo_n] * collectionUuid indexName forward */ @@ -78,6 +80,7 @@ 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, @@ -125,6 +128,7 @@ 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; @@ -141,6 +145,14 @@ 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 263a4f86660..f487c7a45d2 100644 --- a/src/mongo/db/exec/sbe/stages/plan_stats.h +++ b/src/mongo/db/exec/sbe/stages/plan_stats.h @@ -278,8 +278,13 @@ 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}; - long long lastSpilledRecordSize{0}; + // An estimate, in bytes, of the size of the final spill table after all spill events have taken + // place. + long long spilledDataStorageSize{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 678d3f84ef9..3d601d8779e 100644 --- a/src/mongo/db/exec/sbe/stages/scan.cpp +++ b/src/mongo/db/exec/sbe/stages/scan.cpp @@ -209,6 +209,7 @@ void ScanStage::doSaveState(bool relinquishCursor) { cursor->setSaveStorageCursorOnDetachFromOperationContext(!relinquishCursor); } + _indexCatalogEntryMap.clear(); _coll.reset(); } @@ -391,9 +392,14 @@ PlanState ScanStage::getNext() { } // Return EOF if the index key is found to be inconsistent. - if (_scanCallbacks.indexKeyConsistencyCheckCallBack && - !_scanCallbacks.indexKeyConsistencyCheckCallBack( - _opCtx, _snapshotIdAccessor, _indexIdAccessor, _indexKeyAccessor, _coll, *nextRecord)) { + if (_scanCallbacks.indexKeyConsistencyCheckCallback && + !_scanCallbacks.indexKeyConsistencyCheckCallback(_opCtx, + _indexCatalogEntryMap, + _snapshotIdAccessor, + _indexIdAccessor, + _indexKeyAccessor, + _coll, + *nextRecord)) { return trackPlanState(PlanState::IS_EOF); } @@ -457,6 +463,7 @@ void ScanStage::close() { auto optTimer(getOptTimer(_opCtx)); trackClose(); + _indexCatalogEntryMap.clear(); _cursor.reset(); _randomCursor.reset(); _coll.reset(); @@ -742,6 +749,7 @@ void ParallelScanStage::doSaveState(bool relinquishCursor) { _cursor->save(); } + _indexCatalogEntryMap.clear(); _coll.reset(); } @@ -899,8 +907,9 @@ PlanState ParallelScanStage::getNext() { } // Return EOF if the index key is found to be inconsistent. - if (_scanCallbacks.indexKeyConsistencyCheckCallBack && - !_scanCallbacks.indexKeyConsistencyCheckCallBack(_opCtx, + if (_scanCallbacks.indexKeyConsistencyCheckCallback && + !_scanCallbacks.indexKeyConsistencyCheckCallback(_opCtx, + _indexCatalogEntryMap, _snapshotIdAccessor, _indexIdAccessor, _indexKeyAccessor, @@ -956,6 +965,7 @@ 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 37462ac5e14..6f4980f52e1 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|none recordIdSlot|none snapshotIdSlot|none indexIdSlot|none indexKeySlot|none - * indexKeyPatternSlot|none [slot1 = fieldName1, ... slot_n = fieldName_n] collectionUuid + * scan recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? indexKeySlot? + * indexKeyPatternSlot? [slot1 = fieldName1, ... slot_n = fieldName_n] collectionUuid * forward needOplogSlotForTs * - * seek seekKeySlot recordSlot|none recordIdSlot|none snapshotIdSlot|none indexIdSlot|none - * indexKeySlot|none indexKeyPatternSlot|none [slot1 = fieldName1, ... slot_n = fieldName_n] + * seek seekKeySlot recordSlot? recordIdSlot? snapshotIdSlot? indexIdSlot? indexKeySlot? + * indexKeyPatternSlot? [slot1 = fieldName1, ... slot_n = fieldName_n] * collectionUuid forward needOplogSlotForTs */ class ScanStage final : public PlanStage { @@ -197,6 +197,8 @@ 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. @@ -311,6 +313,8 @@ 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 5acf73afe8d..eebdc0c2443 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->numSpills(); + _specificStats.spills += _sorter->stats().spilledRanges(); _specificStats.keysSorted += _sorter->numSorted(); auto& metricsCollector = ResourceConsumption::MetricsCollector::get(_opCtx); metricsCollector.incrementKeysSorted(_sorter->numSorted()); - metricsCollector.incrementSorterSpills(_sorter->numSpills()); + metricsCollector.incrementSorterSpills(_sorter->stats().spilledRanges()); _children[0]->close(); } diff --git a/src/mongo/db/exec/sbe/stages/union.cpp b/src/mongo/db/exec/sbe/stages/union.cpp index a661e6c579f..3ddadb8c912 100644 --- a/src/mongo/db/exec/sbe/stages/union.cpp +++ b/src/mongo/db/exec/sbe/stages/union.cpp @@ -62,28 +62,30 @@ 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]; - auto [it, inserted] = dupCheck.insert(slot); - uassert(4822806, str::stream() << "duplicate field: " << slot, inserted); + bool slotFound = dupCheck.count(slot); + uassert(4822806, str::stream() << "duplicate field: " << slot, !slotFound); 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 c54f3bfe956..6e8c027d7fe 100644 --- a/src/mongo/db/exec/sbe/util/spilling.cpp +++ b/src/mongo/db/exec/sbe/util/spilling.cpp @@ -85,7 +85,6 @@ int upsertToRecordStore(OperationContext* opCtx, 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()); diff --git a/src/mongo/db/exec/sbe/util/spilling.h b/src/mongo/db/exec/sbe/util/spilling.h index 95f73e2b02e..0a562d7e864 100644 --- a/src/mongo/db/exec/sbe/util/spilling.h +++ b/src/mongo/db/exec/sbe/util/spilling.h @@ -55,9 +55,12 @@ boost::optional<value::MaterializedRow> readFromRecordStore(OperationContext* op RecordStore* rs, const RecordId& rid); -/** Inserts or updates a key/value into 'rs'. The 'update' flag controls whether or not an 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, RecordStore* rs, @@ -65,7 +68,6 @@ int upsertToRecordStore(OperationContext* opCtx, const value::MaterializedRow& val, const KeyString::TypeBits& typeBits, bool update); - int upsertToRecordStore(OperationContext* opCtx, RecordStore* rs, const RecordId& key, diff --git a/src/mongo/db/exec/sbe/values/slot.cpp b/src/mongo/db/exec/sbe/values/slot.cpp index 45cbc977980..d8b59c6c3db 100644 --- a/src/mongo/db/exec/sbe/values/slot.cpp +++ b/src/mongo/db/exec/sbe/values/slot.cpp @@ -457,7 +457,7 @@ static void serializeValueIntoKeyString(KeyString::Builder& buf, TypeTags tag, V // TODO SERVER-61629: convert this to serialize the 'arr' directly instead of // constructing a BSONArray. BSONArrayBuilder builder; - bson::convertToBsonObj(builder, getArrayView(val)); + bson::convertToBsonObj(builder, value::ArrayEnumerator{tag, val}); buf.appendBool(true); buf.appendArray(BSONArray(builder.done())); break; @@ -564,8 +564,10 @@ void MaterializedRow::serializeIntoKeyString(KeyString::Builder& buf) const { } } -MaterializedRow MaterializedRow::deserializeFromKeyString(const KeyString::Value& keyString, - BufBuilder* valueBufferBuilder) { +MaterializedRow MaterializedRow::deserializeFromKeyString( + const KeyString::Value& keyString, + BufBuilder* valueBufferBuilder, + boost::optional<size_t> numPrefixValsToRead) { BufReader reader(keyString.getBuffer(), keyString.getSize()); KeyString::TypeBits typeBits(keyString.getTypeBits()); KeyString::TypeBits::Reader typeBitsReader(typeBits); @@ -577,7 +579,8 @@ MaterializedRow MaterializedRow::deserializeFromKeyString(const KeyString::Value &reader, &typeBitsReader, false /* inverted */, typeBits.version, &valBuilder); } while (keepReading); - MaterializedRow result{valBuilder.numValues()}; + size_t sizeOfRow = numPrefixValsToRead ? *numPrefixValsToRead : valBuilder.numValues(); + MaterializedRow result{sizeOfRow}; 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 f853f816d4d..2452c8bd094 100644 --- a/src/mongo/db/exec/sbe/values/slot.h +++ b/src/mongo/db/exec/sbe/values/slot.h @@ -483,10 +483,16 @@ 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. */ - static MaterializedRow deserializeFromKeyString(const KeyString::Value& keyString, + static MaterializedRow deserializeFromKeyString( + const KeyString::Value& keyString, + BufBuilder* valueBufferBuilder, + boost::optional<size_t> numPrefixValsToRead = boost::none); - BufBuilder* valueBufferBuilder); void serializeIntoKeyString(KeyString::Builder& builder) const; private: @@ -577,9 +583,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 { @@ -597,7 +603,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 6c21f4f6e5d..e0f73c87ddf 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()); } -void ArraySet::push_back(TypeTags tag, Value val) { +bool ArraySet::push_back(TypeTags tag, Value val) { if (tag != TypeTags::Nothing) { ValueGuard guard{tag, val}; auto [it, inserted] = _values.insert({tag, val}); @@ -868,7 +868,11 @@ void 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 26ba3e5e1de..ae207106d8b 100644 --- a/src/mongo/db/exec/sbe/values/value.h +++ b/src/mongo/db/exec/sbe/values/value.h @@ -844,7 +844,14 @@ public: } } - void push_back(TypeTags tag, Value val); + /** + * 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); 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 9ad2b511242..00333e9f824 100644 --- a/src/mongo/db/exec/sbe/values/value_builder.h +++ b/src/mongo/db/exec/sbe/values/value_builder.h @@ -191,8 +191,11 @@ 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 < _numValues); + invariant(index < _tagList.size()); auto tag = _tagList[index]; auto val = _valList[index]; @@ -224,9 +227,8 @@ protected: } void appendValue(TypeTags tag, Value val) noexcept { - _tagList[_numValues] = tag; - _valList[_numValues] = val; - ++_numValues; + _tagList.push_back(tag); + _valList.push_back(val); } void appendValue(std::pair<TypeTags, Value> in) noexcept { @@ -241,14 +243,12 @@ 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[_numValues] = tag; - _valList[_numValues] = value::bitcastFrom<int32_t>(_valueBufferBuilder->len()); - ++_numValues; + _tagList.push_back(tag); + _valList.push_back(value::bitcastFrom<int32_t>(_valueBufferBuilder->len())); } - std::array<TypeTags, Ordering::kMaxCompoundIndexKeys> _tagList; - std::array<Value, Ordering::kMaxCompoundIndexKeys> _valList; - size_t _numValues = 0; + absl::InlinedVector<TypeTags, kInlinedVectorSize> _tagList; + absl::InlinedVector<Value, kInlinedVectorSize> _valList; BufBuilder* _valueBufferBuilder; }; @@ -270,11 +270,12 @@ 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. - --_numValues; + _tagList.pop_back(); + _valList.pop_back(); } size_t numValues() const override { - return _numValues; + return _tagList.size(); } /** @@ -284,7 +285,7 @@ public: */ void readValues(std::vector<OwnedValueAccessor>* accessors) { auto bufferLen = _valueBufferBuilder->len(); - for (size_t i = 0; i < _numValues; ++i) { + for (size_t i = 0; i < _tagList.size(); ++i) { auto [tag, val] = getValue(i, bufferLen); invariant(i < accessors->size()); (*accessors)[i].reset(false, tag, val); @@ -304,7 +305,7 @@ public: size_t numValues() const override { size_t nVals = 0; size_t bufIdx = 0; - while (bufIdx < _numValues) { + while (bufIdx < _tagList.size()) { auto tag = _tagList[bufIdx]; auto val = _valList[bufIdx]; if (tag == TypeTags::Boolean && !bitcastTo<bool>(val)) { @@ -323,7 +324,10 @@ public: auto bufferLen = _valueBufferBuilder->len(); size_t bufIdx = 0; size_t rowIdx = 0; - while (bufIdx < _numValues) { + // 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()) { 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 5b44bef6549..3be3212c627 100644 --- a/src/mongo/db/exec/sbe/values/value_serialization_test.cpp +++ b/src/mongo/db/exec/sbe/values/value_serialization_test.cpp @@ -268,6 +268,18 @@ 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)}}); @@ -442,4 +454,17 @@ 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); +} } // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/vm/arith.cpp b/src/mongo/db/exec/sbe/vm/arith.cpp index e4c67775ad8..1d41ff59ce1 100644 --- a/src/mongo/db/exec/sbe/vm/arith.cpp +++ b/src/mongo/db/exec/sbe/vm/arith.cpp @@ -503,6 +503,57 @@ 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; @@ -551,6 +602,67 @@ 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); diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp index 7eb8eb7149e..812de5232af 100644 --- a/src/mongo/db/exec/sbe/vm/vm.cpp +++ b/src/mongo/db/exec/sbe/vm/vm.cpp @@ -1072,35 +1072,40 @@ 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 guard{accTag, accValue}; + value::ValueGuard newArrGuard{accTag, accValue}; auto arr = value::getArrayView(accValue); arr->reserve(AggSumValueElems::kMaxSizeOfArray); - // The order of the following three elements should match to 'AggSumValueElems'. + // 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. 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)); - // 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}; + newArrGuard.reset(); } + + 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}; } @@ -1235,31 +1240,37 @@ 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) { - auto [newAccTag, newAccValue] = value::makeNewArray(); - value::ValueGuard newGuard{newAccTag, newAccValue}; - auto arr = value::getArrayView(newAccValue); + std::tie(accTag, accValue) = value::makeNewArray(); + value::ValueGuard newArrGuard{accTag, accValue}; + auto arr = value::getArrayView(accValue); 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)); - aggStdDevImpl(arr, fieldTag, fieldValue); - newGuard.reset(); - return {true, newAccTag, newAccValue}; + newArrGuard.reset(); } + + 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}; } @@ -3019,6 +3030,120 @@ 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, @@ -3396,6 +3521,166 @@ 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); @@ -4382,7 +4667,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::dispatchBuiltin(Builti case Builtin::doubleDoubleSum: return builtinDoubleDoubleSum(arity); case Builtin::aggDoubleDoubleSum: - return builtinAggDoubleDoubleSum(arity); + return builtinAggDoubleDoubleSum<false /*merging*/>(arity); case Builtin::doubleDoubleSumFinalize: return builtinDoubleDoubleSumFinalize<>(arity); case Builtin::doubleDoubleMergeSumFinalize: @@ -4391,8 +4676,12 @@ 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(arity); + return builtinAggStdDev<false /*merging*/>(arity); + case Builtin::aggMergeStdDevs: + return builtinAggStdDev<true /*merging*/>(arity); case Builtin::stdDevPopFinalize: return builtinStdDevPopFinalize(arity); case Builtin::stdDevSampFinalize: @@ -4445,6 +4734,16 @@ 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 255a1a497c2..e5148e8d4b2 100644 --- a/src/mongo/db/exec/sbe/vm/vm.h +++ b/src/mongo/db/exec/sbe/vm/vm.h @@ -511,12 +511,33 @@ 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 - doubleDoubleSum, // special double summation + + // Special double summation. + doubleDoubleSum, + // A variant of the standard sum aggregate function which maintains a DoubleDouble as the + // accumulator's underlying state. 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 @@ -527,6 +548,18 @@ 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, @@ -980,11 +1013,19 @@ private: value::TypeTags fieldTag, value::Value fieldValue); - void aggDoubleDoubleSumImpl(value::Array* arr, value::TypeTags rhsTag, value::Value rhsValue); + void aggDoubleDoubleSumImpl(value::Array* accumulator, + value::TypeTags rhsTag, + value::Value rhsValue); + void aggMergeDoubleDoubleSumsImpl(value::Array* accumulator, + 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* arr, value::TypeTags rhsTag, value::Value rhsValue); + void aggStdDevImpl(value::Array* accumulator, value::TypeTags rhsTag, value::Value rhsValue); + void aggMergeStdDevsImpl(value::Array* accumulator, + value::TypeTags rhsTag, + value::Value rhsValue); std::tuple<bool, value::TypeTags, value::Value> aggStdDevFinalizeImpl(value::Value fieldValue, bool isSamp); @@ -1103,14 +1144,24 @@ 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); @@ -1137,6 +1188,16 @@ 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); |