SERVER-76952: Skip processing accumulator when the result is no longer affected

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<std::string, BuiltinFn> 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<value::Array*, value::Array*, size_t, size_t, int32_t, int32_t> multi
     return {state, array, startIndexVal, maxSize, memUsage, memLimit};
 }
 
+FastTuple<bool, value::TypeTags, value::Value> 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<size_t>(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<size_t>(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<bool>(needMoreInput)};
+}
+
 int32_t updateAndCheckMemUsage(value::Array* state,
                                int32_t memUsage,
                                int32_t memAdded,
@@ -6122,11 +6146,10 @@ FastTuple<bool, value::TypeTags, value::Value> 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<bool, value::TypeTags, value::Value> 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<bool, value::TypeTags, value::Value> builtinObjectToArray(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinArrayToObject(ArityType arity);
 
+    FastTuple<bool, value::TypeTags, value::Value> builtinAggFirstNNeedsMoreInput(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinAggFirstN(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> builtinAggFirstNMerge(ArityType arity);
     FastTuple<bool, value::TypeTags, value::Value> 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<std::unique_ptr<sbe::EExpression>> buildAccumulatorFirstN(
     boost::optional<sbe::value::SlotId> collatorSlot,
     sbe::value::FrameIdGenerator& frameIdGenerator) {
     std::vector<std::unique_ptr<sbe::EExpression>> aggs;
-    aggs.push_back(makeFunction("aggFirstN", makeFillEmptyNull(std::move(arg))));
+    aggs.push_back(makeLocalBind(
+        &frameIdGenerator,
+        [&](sbe::EVariable accumulatorState) {
+            return sbe::makeE<sbe::EIf>(
+                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) {