SERVER-75486: Support $firstN accumulator in sbe

1 files changed, 228 insertions, 1 deletions

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 03ce1f424eb..d5b52d3489e 100644
--- a/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_accumulator_test.cpp
@@ -31,6 +31,7 @@
 
 #include <fmt/printf.h>
 
+#include "mongo/bson/bsonobj.h"
 #include "mongo/db/exec/sbe/expression_test_base.h"
 #include "mongo/db/pipeline/document_source_group.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
@@ -155,7 +156,7 @@ protected:
                                    ErrorCodes::Error expectedError,
                                    std::unique_ptr<CollatorInterface> collator = nullptr) {
         try {
-            getResultsForAggregation(fromjson(groupSpec.rawData()), inputDocs, std::move(collator));
+            getResultsForAggregation(fromjson(groupSpec), inputDocs, std::move(collator));
             ASSERT(false) << "Expected error: " << expectedError << " for " << groupSpec
                           << " but succeeded";
         } catch (const DBException& e) {
@@ -1605,6 +1606,78 @@ TEST_F(SbeStageBuilderGroupTest, StdDevSampAccumulatorTranslationNonNumber) {
                             BSON_ARRAY(BSON("_id" << BSONNULL << "x" << 0)));
 }
 
+TEST_F(SbeStageBuilderGroupTest, FirstNAccumulatorSingleGroup) {
+    auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 11 << "b" << 1)),
+                                       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))));
+}
+
+TEST_F(SbeStageBuilderGroupTest, FirstNAccumulatorNotEnoughElement) {
+    auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 22 << "b" << 2)),
+                                       BSON_ARRAY(BSON("a" << 11 << "b" << 1))};
+    runGroupAggregationTest("{_id: null, x: {$firstN: {input: '$a', n: 3}}}",
+                            docs,
+                            BSON_ARRAY(BSON("_id" << BSONNULL << "x" << BSON_ARRAY(22 << 11))));
+}
+
+TEST_F(SbeStageBuilderGroupTest, FirstNAccumulatorMultiGroup) {
+    auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 44 << "b" << 4 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 77 << "b" << 7 << "n" << 4)),
+                                       BSON_ARRAY(BSON("a" << 33 << "b" << 3 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 88 << "b" << 8 << "n" << 4)),
+                                       BSON_ARRAY(BSON("a" << 22 << "b" << 2 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 66 << "b" << 6 << "n" << 4)),
+                                       BSON_ARRAY(BSON("a" << 11 << "b" << 1 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 55 << "b" << 5 << "n" << 4))};
+    runGroupAggregationTest("{_id: '$n', x: {$firstN: {input: '$a', n: 3}}}",
+                            docs,
+                            BSON_ARRAY(BSON("_id" << 2 << "x" << BSON_ARRAY(44 << 33 << 22))
+                                       << BSON("_id" << 4 << "x" << BSON_ARRAY(77 << 88 << 66))));
+}
+
+TEST_F(SbeStageBuilderGroupTest, FirstNAccumulatorDynamicN) {
+    auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 44 << "b" << 4 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 33 << "b" << 3 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 22 << "b" << 2 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 11 << "b" << 1 << "n" << 2)),
+                                       BSON_ARRAY(BSON("a" << 88 << "b" << 8 << "n" << 4)),
+                                       BSON_ARRAY(BSON("a" << 77 << "b" << 7 << "n" << 4)),
+                                       BSON_ARRAY(BSON("a" << 66 << "b" << 6 << "n" << 4)),
+                                       BSON_ARRAY(BSON("a" << 55 << "b" << 5 << "n" << 4))};
+    runGroupAggregationTest(
+        "{_id: {k: '$n'}, x: {$firstN: {input: '$a', n: '$k'}}}",
+        docs,
+        BSON_ARRAY(BSON("_id" << BSON("k" << 2) << "x" << BSON_ARRAY(44 << 33))
+                   << BSON("_id" << BSON("k" << 4) << "x" << BSON_ARRAY(88 << 77 << 66 << 55))));
+}
+
+TEST_F(SbeStageBuilderGroupTest, FirstNAccumulatorInvalidConstantN) {
+    const std::vector<std::string> testCases{"'string'", "4.2", "-1", "0"};
+    auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 11 << "b" << 1))};
+    for (const auto& testCase : testCases) {
+        runGroupAggregationToFail(
+            str::stream() << "{_id: null, x: {$firstN: {input: '$a', n: " << testCase << "}}}",
+            docs,
+            static_cast<ErrorCodes::Error>(7548606));
+    }
+}
+
+TEST_F(SbeStageBuilderGroupTest, FirstNAccumulatorInvalidDynamicN) {
+    auto docs = std::vector<BSONArray>{BSON_ARRAY(BSON("a" << 11 << "n" << 1))};
+    runGroupAggregationToFail("{_id: null, x: {$firstN: {input: '$a', n: '$n'}}}",
+                              docs,
+                              static_cast<ErrorCodes::Error>(7548607));
+
+    runGroupAggregationToFail("{_id: {k: '$n'}, x: {$firstN: {input: '$a', n: '$v'}}}",
+                              docs,
+                              static_cast<ErrorCodes::Error>(7548607));
+}
+
 class AccumulatorSBEIncompatible final : public AccumulatorState {
 public:
     static constexpr auto kName = "$incompatible"_sd;
@@ -1941,6 +2014,17 @@ public:
         return {arrTag, arrVal};
     }
 
+    std::pair<sbe::value::TypeTags, sbe::value::Value> convertFromBSONArray(BSONArray arr) {
+        auto [arrTag, arrVal] = sbe::value::makeNewArray();
+        auto arrView = sbe::value::getArrayView(arrVal);
+
+        for (auto elem : arr) {
+            auto [tag, val] = sbe::bson::convertFrom<false>(elem);
+            arrView->push_back(tag, val);
+        }
+        return {arrTag, arrVal};
+    }
+
 protected:
     sbe::value::FrameIdGenerator _frameIdGenerator;
     boost::intrusive_ptr<ExpressionContextForTest> _expCtx;
@@ -2394,4 +2478,147 @@ TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsStdDevSamp) {
     ASSERT_APPROX_EQUAL(sbe::value::bitcastTo<double>(finalizedRes), 3.2660, 0.0001);
 }
 
+TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsFirstNMergeBothArray) {
+    auto expr =
+        stage_builder::makeFunction("aggFirstNMerge", stage_builder::makeVariable(_inputSlotId));
+    auto compiledExpr = compileAggExpression(*expr, &_aggAccessor);
+
+    auto aggSlot = bindAccessor(&_aggAccessor);
+    auto finalizeExpr =
+        stage_builder::makeFunction("aggFirstNFinalize", stage_builder::makeVariable(aggSlot));
+    auto finalizeCompiledExpr = compileExpression(*finalizeExpr);
+
+    // Merge both arrays
+    auto bsonAccArr = BSON_ARRAY(BSON_ARRAY(1 << 2) << 3ll << 16 << 1024);
+    auto [accArrTag, accArrVal] = convertFromBSONArray(bsonAccArr);
+    _aggAccessor.reset(true, accArrTag, accArrVal);
+
+    auto bsonInputArr = BSON_ARRAY(BSON_ARRAY(3 << 4 << 5) << 3ll << 24 << 1024);
+    auto [inputArrTag, inputArrVal] = convertFromBSONArray(bsonInputArr);
+    _inputAccessor.reset(true, inputArrTag, inputArrVal);
+
+    auto [resultTag, resultVal] = runCompiledExpression(compiledExpr.get());
+    _aggAccessor.reset(true, resultTag, resultVal);
+    std::tie(resultTag, resultVal) = runCompiledExpression(finalizeCompiledExpr.get());
+
+    auto expectedArray = BSON_ARRAY(1 << 2 << 3);
+    auto [compareTag, compareVal] =
+        sbe::value::compareValue(resultTag,
+                                 resultVal,
+                                 sbe::value::TypeTags::bsonArray,
+                                 sbe::value::bitcastFrom<const char*>(expectedArray.objdata()));
+
+    ASSERT_EQ(resultTag, sbe::value::TypeTags::Array);
+    ASSERT_EQ(compareTag, sbe::value::TypeTags::NumberInt32);
+    ASSERT_EQ(compareVal, 0);
+    sbe::value::releaseValue(resultTag, resultVal);
+}
+
+TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsFirstNNoMerge) {
+    auto expr =
+        stage_builder::makeFunction("aggFirstNMerge", stage_builder::makeVariable(_inputSlotId));
+    auto compiledExpr = compileAggExpression(*expr, &_aggAccessor);
+
+    auto aggSlot = bindAccessor(&_aggAccessor);
+    auto finalizeExpr =
+        stage_builder::makeFunction("aggFirstNFinalize", stage_builder::makeVariable(aggSlot));
+    auto finalizeCompiledExpr = compileExpression(*finalizeExpr);
+
+    // No merge
+    auto bsonAccArr = BSON_ARRAY(BSON_ARRAY(1 << 2 << 6) << 3ll << 24 << 1024);
+    auto [accArrTag, accArrVal] = convertFromBSONArray(bsonAccArr);
+    _aggAccessor.reset(true, accArrTag, accArrVal);
+
+    auto bsonInputArr = BSON_ARRAY(BSON_ARRAY(3 << 4 << 5) << 3ll << 24 << 1024);
+    auto [inputArrTag, inputArrVal] = convertFromBSONArray(bsonInputArr);
+    _inputAccessor.reset(true, inputArrTag, inputArrVal);
+
+    auto [resultTag, resultVal] = runCompiledExpression(compiledExpr.get());
+    _aggAccessor.reset(true, resultTag, resultVal);
+    std::tie(resultTag, resultVal) = runCompiledExpression(finalizeCompiledExpr.get());
+
+    auto expectedArray = BSON_ARRAY(1 << 2 << 6);
+    auto [compareTag, compareVal] =
+        sbe::value::compareValue(resultTag,
+                                 resultVal,
+                                 sbe::value::TypeTags::bsonArray,
+                                 sbe::value::bitcastFrom<const char*>(expectedArray.objdata()));
+
+    ASSERT_EQ(resultTag, sbe::value::TypeTags::Array);
+    ASSERT_EQ(compareTag, sbe::value::TypeTags::NumberInt32);
+    ASSERT_EQ(compareVal, 0);
+    sbe::value::releaseValue(resultTag, resultVal);
+}
+
+TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsFirstNMergeArrayEmpty) {
+    auto expr =
+        stage_builder::makeFunction("aggFirstNMerge", stage_builder::makeVariable(_inputSlotId));
+    auto compiledExpr = compileAggExpression(*expr, &_aggAccessor);
+
+    auto aggSlot = bindAccessor(&_aggAccessor);
+    auto finalizeExpr =
+        stage_builder::makeFunction("aggFirstNFinalize", stage_builder::makeVariable(aggSlot));
+    auto finalizeCompiledExpr = compileExpression(*finalizeExpr);
+
+    // merge array empty
+    auto bsonAccArr = BSON_ARRAY(BSONArrayBuilder().arr() << 3ll << 0 << 1024);
+    auto [accArrTag, accArrVal] = convertFromBSONArray(bsonAccArr);
+    _aggAccessor.reset(true, accArrTag, accArrVal);
+
+    auto bsonInputArr = BSON_ARRAY(BSON_ARRAY(3 << 4 << 5) << 3ll << 24 << 1024);
+    auto [inputArrTag, inputArrVal] = convertFromBSONArray(bsonInputArr);
+    _inputAccessor.reset(true, inputArrTag, inputArrVal);
+
+    auto [resultTag, resultVal] = runCompiledExpression(compiledExpr.get());
+    _aggAccessor.reset(true, resultTag, resultVal);
+    std::tie(resultTag, resultVal) = runCompiledExpression(finalizeCompiledExpr.get());
+
+    auto expectedArray = BSON_ARRAY(3 << 4 << 5);
+    auto [compareTag, compareVal] =
+        sbe::value::compareValue(resultTag,
+                                 resultVal,
+                                 sbe::value::TypeTags::bsonArray,
+                                 sbe::value::bitcastFrom<const char*>(expectedArray.objdata()));
+
+    ASSERT_EQ(resultTag, sbe::value::TypeTags::Array);
+    ASSERT_EQ(compareTag, sbe::value::TypeTags::NumberInt32);
+    ASSERT_EQ(compareVal, 0);
+    sbe::value::releaseValue(resultTag, resultVal);
+}
+
+TEST_F(SbeStageBuilderGroupAggCombinerTest, CombinePartialAggsFirstNInputArrayEmpty) {
+    auto expr =
+        stage_builder::makeFunction("aggFirstNMerge", stage_builder::makeVariable(_inputSlotId));
+    auto compiledExpr = compileAggExpression(*expr, &_aggAccessor);
+
+    auto aggSlot = bindAccessor(&_aggAccessor);
+    auto finalizeExpr =
+        stage_builder::makeFunction("aggFirstNFinalize", stage_builder::makeVariable(aggSlot));
+    auto finalizeCompiledExpr = compileExpression(*finalizeExpr);
+
+    // input array empty
+    auto bsonAccArr = BSON_ARRAY(BSON_ARRAY(3 << 4 << 5) << 3ll << 24 << 1024);
+    auto [accArrTag, accArrVal] = convertFromBSONArray(bsonAccArr);
+    _aggAccessor.reset(true, accArrTag, accArrVal);
+
+    auto bsonInputArr = BSON_ARRAY(BSONArrayBuilder().arr() << 3ll << 0 << 1024);
+    auto [inputArrTag, inputArrVal] = convertFromBSONArray(bsonInputArr);
+    _inputAccessor.reset(true, inputArrTag, inputArrVal);
+
+    auto [resultTag, resultVal] = runCompiledExpression(compiledExpr.get());
+    _aggAccessor.reset(true, resultTag, resultVal);
+    std::tie(resultTag, resultVal) = runCompiledExpression(finalizeCompiledExpr.get());
+
+    auto expectedArray = BSON_ARRAY(3 << 4 << 5);
+    auto [compareTag, compareVal] =
+        sbe::value::compareValue(resultTag,
+                                 resultVal,
+                                 sbe::value::TypeTags::bsonArray,
+                                 sbe::value::bitcastFrom<const char*>(expectedArray.objdata()));
+
+    ASSERT_EQ(resultTag, sbe::value::TypeTags::Array);
+    ASSERT_EQ(compareTag, sbe::value::TypeTags::NumberInt32);
+    ASSERT_EQ(compareVal, 0);
+    sbe::value::releaseValue(resultTag, resultVal);
+}
 }  // namespace mongo