From 427574c5614a3a3587b4a6563823149d60e71543 Mon Sep 17 00:00:00 2001 From: Projjal Chanda Date: Mon, 15 May 2023 17:39:03 +0000 Subject: SERVER-76952: Skip processing accumulator when the result is no longer affected --- src/mongo/db/exec/sbe/expressions/expression.cpp | 16 ++++++++++- src/mongo/db/exec/sbe/vm/vm.cpp | 33 ++++++++++++++++++++-- src/mongo/db/exec/sbe/vm/vm.h | 4 ++- .../db/query/sbe_stage_builder_accumulator.cpp | 13 ++++++++- .../query/sbe_stage_builder_accumulator_test.cpp | 20 +++++++------ 5 files changed, 72 insertions(+), 14 deletions(-) diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp index 5f2c160b48e..67fb7d8c579 100644 --- a/src/mongo/db/exec/sbe/expressions/expression.cpp +++ b/src/mongo/db/exec/sbe/expressions/expression.cpp @@ -782,7 +782,9 @@ static stdx::unordered_map kBuiltinFunctions = { {"arrayToObject", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::arrayToObject, false}}, {"array", BuiltinFn{kAnyNumberOfArgs, vm::Builtin::newArray, false}}, - {"aggFirstN", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggFirstN, true}}, + {"aggFirstNNeedsMoreInput", + BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggFirstNNeedsMoreInput, false}}, + {"aggFirstN", BuiltinFn{[](size_t n) { return n == 2; }, vm::Builtin::aggFirstN, false}}, {"aggFirstNMerge", BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::aggFirstNMerge, true}}, {"aggFirstNFinalize", @@ -1150,6 +1152,18 @@ vm::CodeFragment EFunction::compileDirect(CompileCtx& ctx) const { return (it->second.generate)(ctx, _nodes, it->second.aggregate); } + if (_name == "aggState") { + uassert(7695204, + "aggregate function call: aggState occurs in the non-aggregate context.", + ctx.aggExpression && ctx.accumulator); + uassert(7695205, + str::stream() << "function call: aggState has wrong arity: " << _nodes.size(), + _nodes.size() == 0); + vm::CodeFragment code; + code.appendMoveVal(ctx.accumulator); + return code; + } + uasserted(4822847, str::stream() << "unknown function call: " << _name); } diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp index 2f7c1bee391..f79ea074019 100644 --- a/src/mongo/db/exec/sbe/vm/vm.cpp +++ b/src/mongo/db/exec/sbe/vm/vm.cpp @@ -6076,6 +6076,30 @@ std::tuple multi return {state, array, startIndexVal, maxSize, memUsage, memLimit}; } +FastTuple ByteCode::builtinAggFirstNNeedsMoreInput( + ArityType arity) { + auto [stateOwned, stateTag, stateVal] = getFromStack(0); + uassert(7695200, "Unexpected accumulator state ownership", !stateOwned); + + auto state = value::getArrayView(stateVal); + uassert( + 7695201, "The accumulator state should be an array", stateTag == value::TypeTags::Array); + + auto [arrayTag, arrayVal] = state->getAt(static_cast(AggMultiElems::kInternalArr)); + uassert(7695202, + "Internal array component is not of correct type", + arrayTag == value::TypeTags::Array); + auto array = value::getArrayView(arrayVal); + + auto [maxSizeTag, maxSize] = state->getAt(static_cast(AggMultiElems::kMaxSize)); + uassert(7695203, + "MaxSize component should be a 64-bit integer", + maxSizeTag == value::TypeTags::NumberInt64); + + bool needMoreInput = (array->size() < maxSize); + return {false, value::TypeTags::Boolean, value::bitcastFrom(needMoreInput)}; +} + int32_t updateAndCheckMemUsage(value::Array* state, int32_t memUsage, int32_t memAdded, @@ -6122,11 +6146,10 @@ FastTuple ByteCode::builtinAggFirstN(ArityT auto [stateTag, stateVal] = moveOwnedFromStack(0); value::ValueGuard stateGuard{stateTag, stateVal}; - auto [state, array, accStartIdx, accSize, memUsage, memLimit] = - multiAccState(stateTag, stateVal); + auto [state, array, startIdx, maxSize, memUsage, memLimit] = multiAccState(stateTag, stateVal); auto [fieldTag, fieldVal] = moveOwnedFromStack(1); - aggFirstN(state, array, accSize, memUsage, memLimit, fieldTag, fieldVal); + aggFirstN(state, array, maxSize, memUsage, memLimit, fieldTag, fieldVal); stateGuard.reset(); return {true, stateTag, stateVal}; @@ -6813,6 +6836,8 @@ FastTuple ByteCode::dispatchBuiltin(Builtin return builtinObjectToArray(arity); case Builtin::arrayToObject: return builtinArrayToObject(arity); + case Builtin::aggFirstNNeedsMoreInput: + return builtinAggFirstNNeedsMoreInput(arity); case Builtin::aggFirstN: return builtinAggFirstN(arity); case Builtin::aggFirstNMerge: @@ -7121,6 +7146,8 @@ std::string builtinToString(Builtin b) { return "objectToArray"; case Builtin::arrayToObject: return "arrayToObject"; + case Builtin::aggFirstNNeedsMoreInput: + return "aggFirstNNeedsMoreInput"; case Builtin::aggFirstN: return "aggFirstN"; case Builtin::aggFirstNMerge: diff --git a/src/mongo/db/exec/sbe/vm/vm.h b/src/mongo/db/exec/sbe/vm/vm.h index 17b3eb29b7f..07cf0d20102 100644 --- a/src/mongo/db/exec/sbe/vm/vm.h +++ b/src/mongo/db/exec/sbe/vm/vm.h @@ -761,6 +761,7 @@ enum class Builtin : uint8_t { objectToArray, arrayToObject, + aggFirstNNeedsMoreInput, aggFirstN, aggFirstNMerge, aggFirstNFinalize, @@ -786,7 +787,7 @@ std::string builtinToString(Builtin b); /** * This enum defines indices into an 'Array' that store state for $AccumulatorN expressions. * - * The array might contain up to four elements: + * The array contains five elements: * - The element at index `kInternalArr` is the array that holds the values. * - The element at index `kStartIdx` is the logical start index in the internal array. This is * used for emulating queue behaviour. @@ -1691,6 +1692,7 @@ private: FastTuple builtinObjectToArray(ArityType arity); FastTuple builtinArrayToObject(ArityType arity); + FastTuple builtinAggFirstNNeedsMoreInput(ArityType arity); FastTuple builtinAggFirstN(ArityType arity); FastTuple builtinAggFirstNMerge(ArityType arity); FastTuple builtinAggFirstNFinalize(ArityType arity); diff --git a/src/mongo/db/query/sbe_stage_builder_accumulator.cpp b/src/mongo/db/query/sbe_stage_builder_accumulator.cpp index 2e060e2258c..279b1071f5e 100644 --- a/src/mongo/db/query/sbe_stage_builder_accumulator.cpp +++ b/src/mongo/db/query/sbe_stage_builder_accumulator.cpp @@ -651,7 +651,18 @@ std::vector> buildAccumulatorFirstN( boost::optional collatorSlot, sbe::value::FrameIdGenerator& frameIdGenerator) { std::vector> aggs; - aggs.push_back(makeFunction("aggFirstN", makeFillEmptyNull(std::move(arg)))); + aggs.push_back(makeLocalBind( + &frameIdGenerator, + [&](sbe::EVariable accumulatorState) { + return sbe::makeE( + makeFunction("aggFirstNNeedsMoreInput", accumulatorState.clone()), + makeFunction( + "aggFirstN", + makeMoveVariable(*accumulatorState.getFrameId(), accumulatorState.getSlotId()), + makeFillEmptyNull(std::move(arg))), + makeMoveVariable(*accumulatorState.getFrameId(), accumulatorState.getSlotId())); + }, + makeFunction("aggState"))); return aggs; } diff --git a/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp b/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp index 7de06bf4e4f..f4fb59f42f5 100644 --- a/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp +++ b/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp @@ -1612,15 +1612,19 @@ TEST_F(SbeStageBuilderGroupTest, FirstNLastNAccumulatorSingleGroup) { BSON_ARRAY(BSON("a" << 22 << "b" << 2)), BSON_ARRAY(BSON("a" << 33 << "b" << 3)), BSON_ARRAY(BSON("a" << 44 << "b" << 4))}; - runGroupAggregationTest( - "{_id: null, x: {$firstN: {input: '$a', n: 3}}}", - docs, - BSON_ARRAY(BSON("_id" << BSONNULL << "x" << BSON_ARRAY(11 << 22 << 33)))); + runGroupAggregationTest("{_id: null, x: {$firstN: {input: {a: '$a', b: '$b'}, n: 3}}}", + docs, + BSON_ARRAY(BSON("_id" << BSONNULL << "x" + << BSON_ARRAY(BSON("a" << 11 << "b" << 1) + << BSON("a" << 22 << "b" << 2) + << BSON("a" << 33 << "b" << 3))))); - runGroupAggregationTest( - "{_id: null, x: {$lastN: {input: '$a', n: 3}}}", - docs, - BSON_ARRAY(BSON("_id" << BSONNULL << "x" << BSON_ARRAY(22 << 33 << 44)))); + runGroupAggregationTest("{_id: null, x: {$lastN: {input: {a: '$a', b: '$b'}, n: 3}}}", + docs, + BSON_ARRAY(BSON("_id" << BSONNULL << "x" + << BSON_ARRAY(BSON("a" << 22 << "b" << 2) + << BSON("a" << 33 << "b" << 3) + << BSON("a" << 44 << "b" << 4))))); } TEST_F(SbeStageBuilderGroupTest, FirstNLastNAccumulatorNotEnoughElement) { -- cgit v1.2.1