SERVER-51267 Support $arrayElemAt, $first and $last expressions in SBE

12 files changed, 443 insertions, 126 deletions

diff --git a/jstests/aggregation/bugs/server4589.js b/jstests/aggregation/bugs/server4589.js
index 623120e631b..59de58be383 100644
--- a/jstests/aggregation/bugs/server4589.js
+++ b/jstests/aggregation/bugs/server4589.js
@@ -1,10 +1,7 @@
 // SERVER-4589: Add $arrayElemAt aggregation expression.
-// @tags: [
-//   sbe_incompatible,
-// ]
 
-// For assertErrorCode.
-load('jstests/aggregation/extras/utils.js');
+load('jstests/aggregation/extras/utils.js');        // For assertErrorCode.
+load("jstests/libs/sbe_assert_error_override.js");  // Override error-code-checking APIs.
 
 (function() {
 'use strict';
diff --git a/jstests/aggregation/expressions/first_last.js b/jstests/aggregation/expressions/first_last.js
index 501a3dc2324..f0edd042cb6 100644
--- a/jstests/aggregation/expressions/first_last.js
+++ b/jstests/aggregation/expressions/first_last.js
@@ -18,20 +18,20 @@ assert.commandWorked(coll.insert([
 ]));
 
 const result =
-    coll.aggregate([{$sort: {_id: 1}}, {$addFields: {f: {$first: "$a"}, l: {$last: "$a"}}}])
+    coll.aggregate([{$sort: {_id: 1}}, {$project: {f: {$first: "$a"}, l: {$last: "$a"}}}])
         .toArray();
 assert.eq(result, [
     // When an array doesn't contain a given index, the result is 'missing', similar to looking up a
     // nonexistent key in a document.
-    {_id: 0, a: []},
+    {_id: 0},
 
-    {_id: 1, a: ['A'], f: 'A', l: 'A'},
-    {_id: 2, a: ['A', 'B'], f: 'A', l: 'B'},
-    {_id: 3, a: ['A', 'B', 'C'], f: 'A', l: 'C'},
+    {_id: 1, f: 'A', l: 'A'},
+    {_id: 2, f: 'A', l: 'B'},
+    {_id: 3, f: 'A', l: 'C'},
 
     // When the input is nullish instead of an array, the result is null.
-    {_id: 4, a: null, f: null, l: null},
-    {_id: 5, a: undefined, f: null, l: null},
+    {_id: 4, f: null, l: null},
+    {_id: 5, f: null, l: null},
     {_id: 6, f: null, l: null},
 ]);
 }());
diff --git a/jstests/libs/sbe_assert_error_override.js b/jstests/libs/sbe_assert_error_override.js
index ed6fa7a0b26..8d6147bccd9 100644
--- a/jstests/libs/sbe_assert_error_override.js
+++ b/jstests/libs/sbe_assert_error_override.js
@@ -23,12 +23,18 @@
 const equivalentErrorCodesList = [
     [28651, 5073201],
     [16006, 4997703],
+    [28689, 5126701],
+    [28690, 5126702],
+    [28691, 5126703],
     [16020, 5066300],
     [16007, 5066300],
     [16608, 4848401],
     [16609, 5073101],
     [16555, 5073102],
     [28680, 4903701],
+    [28689, 5126701],
+    [28690, 5126702],
+    [28691, 5126703],
     [28765, 4822870],
     [28714, 4903710],
     [28761, 4903708],
diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index 787797e83c2..3942f120a05 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -96,6 +96,7 @@ env.CppUnitTest(
         'expressions/sbe_index_of_test.cpp',
         'expressions/sbe_to_upper_to_lower_test.cpp',
         'expressions/sbe_trigonometric_expressions_test.cpp',
+        'expressions/sbe_get_element_builtin_test.cpp',
         'sbe_filter_test.cpp',
         'sbe_key_string_test.cpp',
         'sbe_limit_skip_test.cpp',
diff --git a/src/mongo/db/exec/sbe/expression_test_base.h b/src/mongo/db/exec/sbe/expression_test_base.h
index 2ce40a79741..36badd2d5e0 100644
--- a/src/mongo/db/exec/sbe/expression_test_base.h
+++ b/src/mongo/db/exec/sbe/expression_test_base.h
@@ -90,6 +90,75 @@ protected:
         return _vm.runPredicate(compiledExpr);
     }
 
+    std::pair<value::TypeTags, value::Value> makeBsonArray(const BSONArray& ba) {
+        int numBytes = ba.objsize();
+        uint8_t* data = new uint8_t[numBytes];
+        memcpy(data, reinterpret_cast<const uint8_t*>(ba.objdata()), numBytes);
+        return {value::TypeTags::bsonArray, value::bitcastFrom<uint8_t*>(data)};
+    }
+
+    std::pair<value::TypeTags, value::Value> makeArraySet(const BSONArray& arr) {
+        auto [tmpTag, tmpVal] = makeBsonArray(arr);
+        value::ValueGuard tmpGuard{tmpTag, tmpVal};
+
+        value::ArrayEnumerator enumerator{tmpTag, tmpVal};
+
+        auto [arrTag, arrVal] = value::makeNewArraySet();
+        value::ValueGuard guard{arrTag, arrVal};
+
+        auto arrView = value::getArraySetView(arrVal);
+
+        while (!enumerator.atEnd()) {
+            auto [tag, val] = enumerator.getViewOfValue();
+            enumerator.advance();
+
+            auto [copyTag, copyVal] = value::copyValue(tag, val);
+            arrView->push_back(copyTag, copyVal);
+        }
+        guard.reset();
+
+        return {arrTag, arrVal};
+    }
+
+    std::pair<value::TypeTags, value::Value> makeArray(const BSONArray& arr) {
+        auto [tmpTag, tmpVal] = makeBsonArray(arr);
+        value::ValueGuard tmpGuard{tmpTag, tmpVal};
+
+        value::ArrayEnumerator enumerator{tmpTag, tmpVal};
+
+        auto [arrTag, arrVal] = value::makeNewArray();
+        value::ValueGuard guard{arrTag, arrVal};
+
+        auto arrView = value::getArrayView(arrVal);
+
+        while (!enumerator.atEnd()) {
+            auto [tag, val] = enumerator.getViewOfValue();
+            enumerator.advance();
+
+            auto [copyTag, copyVal] = value::copyValue(tag, val);
+            arrView->push_back(copyTag, copyVal);
+        }
+        guard.reset();
+
+        return {arrTag, arrVal};
+    }
+
+    std::pair<value::TypeTags, value::Value> makeNothing() {
+        return {value::TypeTags::Nothing, value::bitcastFrom<int64_t>(0)};
+    }
+
+    std::pair<value::TypeTags, value::Value> makeInt32(int32_t value) {
+        return {value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(value)};
+    }
+
+    std::pair<value::TypeTags, value::Value> makeInt64(int64_t value) {
+        return {value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(value)};
+    }
+
+    std::pair<value::TypeTags, value::Value> makeDouble(double value) {
+        return {value::TypeTags::NumberDouble, value::bitcastFrom<double>(value)};
+    }
+
 private:
     value::SlotIdGenerator _slotIdGenerator;
     CoScanStage _emptyStage{kEmptyPlanNodeId};
diff --git a/src/mongo/db/exec/sbe/expressions/sbe_get_element_builtin_test.cpp b/src/mongo/db/exec/sbe/expressions/sbe_get_element_builtin_test.cpp
new file mode 100644
index 00000000000..a6b1c348781
--- /dev/null
+++ b/src/mongo/db/exec/sbe/expressions/sbe_get_element_builtin_test.cpp
@@ -0,0 +1,175 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include <cmath>
+#include <limits>
+
+#include "mongo/db/exec/sbe/expression_test_base.h"
+#include "mongo/db/exec/sbe/values/bson.h"
+
+namespace mongo::sbe {
+
+class SBEBuiltinGetElementTest : public EExpressionTestFixture {
+protected:
+    using TypedValue = std::pair<value::TypeTags, value::Value>;
+
+    struct TestCase {
+        BSONArray array;
+        TypedValue index;
+        TypedValue expected;
+    };
+
+    void setUp() override {
+        decimalValue = value::makeCopyDecimal(Decimal128("1.2345"));
+        testCases = {
+            // Positive indexes.
+            {BSON_ARRAY(1 << 2 << 3), makeInt32(0), makeInt32(1)},
+            {BSON_ARRAY(1 << 2 << 3), makeInt32(1), makeInt32(2)},
+            {BSON_ARRAY(1 << 2 << 3), makeInt32(2), makeInt32(3)},
+
+            // Negative indexes.
+            {BSON_ARRAY(1 << 2 << 3), makeInt32(-1), makeInt32(3)},
+            {BSON_ARRAY(1 << 2 << 3), makeInt32(-2), makeInt32(2)},
+            {BSON_ARRAY(1 << 2 << 3), makeInt32(-3), makeInt32(1)},
+
+            // Index out of bounds.
+            {BSON_ARRAY(1 << 2), makeInt32(2), makeNothing()},
+            {BSON_ARRAY(1 << 2), makeInt32(-3), makeNothing()},
+            {BSONArray(), makeInt32(0), makeNothing()},
+            {BSONArray(), makeInt32(-1), makeNothing()},
+
+            // Invalid index type.
+            {BSON_ARRAY(1 << 2), makeNothing(), makeNothing()},
+            {BSON_ARRAY(1 << 2), decimalValue, makeNothing()},
+            {BSON_ARRAY(1 << 2), makeDouble(1.2345), makeNothing()},
+            {BSON_ARRAY(1 << 2), makeInt64(0), makeNothing()},
+        };
+    }
+
+    void tearDown() override {
+        value::releaseValue(decimalValue.first, decimalValue.second);
+    }
+
+    /**
+     * Compile and run expression 'getElement(array, index)' and return its result.
+     * NOTE: Values behind arguments and the return value of this function are owned by the caller.
+     */
+    TypedValue runExpression(TypedValue array, TypedValue index) {
+        // We do not copy array value on purpose. During the copy, order of elements in ArraySet
+        // may change. 'getElement' return value depends on the order of elements in the input
+        // array. After copy 'getElement' may return different element from the expected by the
+        // caller.
+        value::ViewOfValueAccessor arraySlotAccessor;
+        auto arraySlot = bindAccessor(&arraySlotAccessor);
+        arraySlotAccessor.reset(array.first, array.second);
+        auto arrayExpr = makeE<EVariable>(arraySlot);
+
+        auto indexCopy = value::copyValue(index.first, index.second);
+        auto indexExpr = makeE<EConstant>(indexCopy.first, indexCopy.second);
+
+        auto getElementExpr =
+            makeE<EFunction>("getElement", makeEs(std::move(arrayExpr), std::move(indexExpr)));
+        auto compiledExpr = compileExpression(*getElementExpr);
+
+        return runCompiledExpression(compiledExpr.get());
+    }
+
+    /**
+     * Assert that result of 'getElement(array, index)' is equal to 'expectedRes'.
+     */
+    void runAndAssertExpression(TypedValue array, TypedValue index, TypedValue expectedRes) {
+        auto actualValue = runExpression(array, index);
+        value::ValueGuard guard{actualValue};
+
+        auto [compareTag, compareValue] = value::compareValue(
+            actualValue.first, actualValue.second, expectedRes.first, expectedRes.second);
+        ASSERT_EQ(compareTag, value::TypeTags::NumberInt32);
+        ASSERT_EQ(compareValue, 0);
+    }
+
+    TypedValue decimalValue;
+    std::vector<TestCase> testCases;
+};
+
+TEST_F(SBEBuiltinGetElementTest, GetElementBSONArray) {
+    for (const auto& testCase : testCases) {
+        auto bsonArray = makeBsonArray(testCase.array);
+        value::ValueGuard guard{bsonArray};
+        runAndAssertExpression(bsonArray, testCase.index, testCase.expected);
+    }
+}
+
+TEST_F(SBEBuiltinGetElementTest, GetElementArray) {
+    for (const auto& testCase : testCases) {
+        auto array = makeArray(testCase.array);
+        value::ValueGuard guard{array};
+        runAndAssertExpression(array, testCase.index, testCase.expected);
+    }
+}
+
+TEST_F(SBEBuiltinGetElementTest, GetElementArraySetNothing) {
+    // Run test cases when 'getElement' returns Nothing.
+    for (const auto& testCase : testCases) {
+        if (testCase.expected.first != value::TypeTags::Nothing) {
+            continue;
+        }
+
+        auto array = makeArraySet(testCase.array);
+        value::ValueGuard guard{array};
+        runAndAssertExpression(array, testCase.index, testCase.expected);
+    }
+}
+
+TEST_F(SBEBuiltinGetElementTest, GetElementArraySetElements) {
+    auto array = makeArraySet(BSON_ARRAY(1 << 2 << 3));
+    value::ValueGuard guard{array};
+
+    const std::vector<std::pair<int32_t, int32_t>> indices = {{-3, -1}, {0, 2}};
+    for (const auto& [begin, end] : indices) {
+        std::vector<int32_t> elements;
+        for (int32_t i = begin; i <= end; ++i) {
+            auto result = runExpression(array, makeInt32(i));
+            value::ValueGuard guard{result};
+            ASSERT_EQ(result.first, value::TypeTags::NumberInt32);
+            elements.push_back(value::bitcastTo<int32_t>(result.second));
+        }
+
+        std::sort(elements.begin(), elements.end());
+        ASSERT_EQ(elements[0], 1);
+        ASSERT_EQ(elements[1], 2);
+        ASSERT_EQ(elements[2], 3);
+    }
+}
+
+TEST_F(SBEBuiltinGetElementTest, GetElementNotArray) {
+    runAndAssertExpression(makeNothing(), makeInt32(1), makeNothing());
+    runAndAssertExpression(makeInt32(123), makeInt32(1), makeNothing());
+}
+
+}  // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/expressions/sbe_is_member_builtin_test.cpp b/src/mongo/db/exec/sbe/expressions/sbe_is_member_builtin_test.cpp
index 80b45d16505..fd0781606ba 100644
--- a/src/mongo/db/exec/sbe/expressions/sbe_is_member_builtin_test.cpp
+++ b/src/mongo/db/exec/sbe/expressions/sbe_is_member_builtin_test.cpp
@@ -45,62 +45,9 @@ protected:
         ASSERT_EQ(actualRes, expectedRes);
     }
 
-    std::pair<value::TypeTags, value::Value> makeValue(const BSONArray& ba) {
-        int numBytes = ba.objsize();
-        uint8_t* data = new uint8_t[numBytes];
-        memcpy(data, reinterpret_cast<const uint8_t*>(ba.objdata()), numBytes);
-        return {value::TypeTags::bsonArray, value::bitcastFrom<uint8_t*>(data)};
-    }
-
     std::pair<value::TypeTags, value::Value> makeViewOfObject(const BSONObj& obj) {
         return {value::TypeTags::bsonObject, value::bitcastFrom<const char*>(obj.objdata())};
     }
-
-    std::pair<value::TypeTags, value::Value> makeArraySet(const BSONArray& arr) {
-        auto [tmpTag, tmpVal] = makeValue(arr);
-        value::ValueGuard tmpGuard{tmpTag, tmpVal};
-
-        value::ArrayEnumerator enumerator{tmpTag, tmpVal};
-
-        auto [arrTag, arrVal] = value::makeNewArraySet();
-        value::ValueGuard guard{arrTag, arrVal};
-
-        auto arrView = value::getArraySetView(arrVal);
-
-        while (!enumerator.atEnd()) {
-            auto [tag, val] = enumerator.getViewOfValue();
-            enumerator.advance();
-
-            auto [copyTag, copyVal] = value::copyValue(tag, val);
-            arrView->push_back(copyTag, copyVal);
-        }
-        guard.reset();
-
-        return {arrTag, arrVal};
-    }
-
-    std::pair<value::TypeTags, value::Value> makeArray(const BSONArray& arr) {
-        auto [tmpTag, tmpVal] = makeValue(arr);
-        value::ValueGuard tmpGuard{tmpTag, tmpVal};
-
-        value::ArrayEnumerator enumerator{tmpTag, tmpVal};
-
-        auto [arrTag, arrVal] = value::makeNewArray();
-        value::ValueGuard guard{arrTag, arrVal};
-
-        auto arrView = value::getArrayView(arrVal);
-
-        while (!enumerator.atEnd()) {
-            auto [tag, val] = enumerator.getViewOfValue();
-            enumerator.advance();
-
-            auto [copyTag, copyVal] = value::copyValue(tag, val);
-            arrView->push_back(copyTag, copyVal);
-        }
-        guard.reset();
-
-        return {arrTag, arrVal};
-    }
 };
 
 TEST_F(SBEBuiltinIsMemberTest, IsMemberArraySet) {
@@ -238,50 +185,50 @@ TEST_F(SBEBuiltinIsMemberTest, IsMemberBSONArray) {
 
     // Test that isMember can find basic values.
     inputSlotAccessor.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(1));
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(1 << 2)), true);
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(1 << 2)), true);
 
     inputSlotAccessor.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(3));
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(1 << 2)), false);
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(1 << 2)), false);
 
     inputSlotAccessor.reset(value::TypeTags::NumberInt32, value::bitcastFrom<int32_t>(3));
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(BSONObj())), false);
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(BSONObj())), false);
 
-    inputSlotAccessor.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1));
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(1 << 2)), true);
+    inputSlotAccessor.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int32_t>(1));
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(1 << 2)), true);
 
-    inputSlotAccessor.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(3));
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(1 << 2)), false);
+    inputSlotAccessor.reset(value::TypeTags::NumberInt64, value::bitcastFrom<int32_t>(3));
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(1 << 2)), false);
 
     auto [decimalTag, decimalVal] = value::makeCopyDecimal(Decimal128{9});
     inputSlotAccessor.reset(decimalTag, decimalVal);
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(1 << 9.0)), true);
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(1 << 9.0)), true);
 
     std::tie(decimalTag, decimalVal) = value::makeCopyDecimal(Decimal128{0.1});
     inputSlotAccessor.reset(decimalTag, decimalVal);
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(1 << 9.0)), false);
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(1 << 9.0)), false);
 
     auto [smallStrTag, smallStrVal] = value::makeSmallString("foo");
     inputSlotAccessor.reset(smallStrTag, smallStrVal);
     runAndAssertExpression(inputSlot,
-                           makeValue(BSON_ARRAY("foo"
-                                                << "bar")),
+                           makeBsonArray(BSON_ARRAY("foo"
+                                                    << "bar")),
                            true);
 
     std::tie(smallStrTag, smallStrVal) = value::makeSmallString("baz");
     inputSlotAccessor.reset(smallStrTag, smallStrVal);
     runAndAssertExpression(inputSlot,
-                           makeValue(BSON_ARRAY("foo"
-                                                << "bar")),
+                           makeBsonArray(BSON_ARRAY("foo"
+                                                    << "bar")),
                            false);
 
     // Test that isMember can find composite values.
     auto [arrTag, arrVal] = makeArray(BSON_ARRAY(2 << 3));
     inputSlotAccessor.reset(arrTag, arrVal);
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(BSON_ARRAY(2 << 3) << 1)), true);
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(BSON_ARRAY(2 << 3) << 1)), true);
 
     std::tie(arrTag, arrVal) = makeArray(BSON_ARRAY(1));
     inputSlotAccessor.reset(arrTag, arrVal);
-    runAndAssertExpression(inputSlot, makeValue(BSON_ARRAY(BSON_ARRAY(1 << 2) << 3)), false);
+    runAndAssertExpression(inputSlot, makeBsonArray(BSON_ARRAY(BSON_ARRAY(1 << 2) << 3)), false);
 
     value::ViewOfValueAccessor bsonObjAccessor;
     auto bsonObjSlot = bindAccessor(&bsonObjAccessor);
@@ -290,10 +237,10 @@ TEST_F(SBEBuiltinIsMemberTest, IsMemberBSONArray) {
     auto [targetTag, targetVal] = makeViewOfObject(targetObj);
 
     bsonObjAccessor.reset(targetTag, targetVal);
-    runAndAssertExpression(bsonObjSlot, makeValue(BSON_ARRAY(1 << BSON("a" << 1))), true);
+    runAndAssertExpression(bsonObjSlot, makeBsonArray(BSON_ARRAY(1 << BSON("a" << 1))), true);
 
     bsonObjAccessor.reset(targetTag, targetVal);
-    runAndAssertExpression(bsonObjSlot, makeValue(BSON_ARRAY(10 << BSON("b" << 1))), false);
+    runAndAssertExpression(bsonObjSlot, makeBsonArray(BSON_ARRAY(10 << BSON("b" << 1))), false);
 }
 
 
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
index e6529e317e8..c634160f3d3 100644
--- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.cpp
@@ -83,7 +83,7 @@ std::pair<value::SlotId, std::unique_ptr<PlanStage>> PlanStageTestFixture::gener
 }
 
 std::pair<value::SlotVector, std::unique_ptr<PlanStage>>
-PlanStageTestFixture::generateMockScanMulti(int64_t numSlots,
+PlanStageTestFixture::generateMockScanMulti(int32_t numSlots,
                                             value::TypeTags arrTag,
                                             value::Value arrVal) {
     using namespace std::literals;
@@ -97,14 +97,14 @@ PlanStageTestFixture::generateMockScanMulti(int64_t numSlots,
     // across multiple output slots.
     value::SlotVector projectSlots;
     value::SlotMap<std::unique_ptr<EExpression>> projections;
-    for (int64_t i = 0; i < numSlots; ++i) {
+    for (int32_t i = 0; i < numSlots; ++i) {
         projectSlots.emplace_back(generateSlotId());
         projections.emplace(
             projectSlots.back(),
             makeE<EFunction>("getElement"sv,
                              makeEs(makeE<EVariable>(scanSlot),
-                                    makeE<EConstant>(value::TypeTags::NumberInt64,
-                                                     value::bitcastFrom<int64_t>(i)))));
+                                    makeE<EConstant>(value::TypeTags::NumberInt32,
+                                                     value::bitcastFrom<int32_t>(i)))));
     }
 
     return {std::move(projectSlots),
@@ -118,7 +118,7 @@ std::pair<value::SlotId, std::unique_ptr<PlanStage>> PlanStageTestFixture::gener
 }
 
 std::pair<value::SlotVector, std::unique_ptr<PlanStage>>
-PlanStageTestFixture::generateMockScanMulti(int64_t numSlots, const BSONArray& array) {
+PlanStageTestFixture::generateMockScanMulti(int32_t numSlots, const BSONArray& array) {
     auto [arrTag, arrVal] = makeValue(array);
     return generateMockScanMulti(numSlots, arrTag, arrVal);
 }
@@ -214,7 +214,7 @@ void PlanStageTestFixture::runTest(value::TypeTags inputTag,
     ASSERT_TRUE(valueEquals(resultsTag, resultsVal, expectedTag, expectedVal));
 }
 
-void PlanStageTestFixture::runTestMulti(int64_t numInputSlots,
+void PlanStageTestFixture::runTestMulti(int32_t numInputSlots,
                                         value::TypeTags inputTag,
                                         value::Value inputVal,
                                         value::TypeTags expectedTag,
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
index ce636d4f7ad..c8d74595c35 100644
--- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
+++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
@@ -152,7 +152,7 @@ public:
      * Note that this method assumes ownership of the SBE Array being passed in.
      */
     std::pair<value::SlotVector, std::unique_ptr<PlanStage>> generateMockScanMulti(
-        int64_t numSlots, value::TypeTags arrTag, value::Value arrVal);
+        int32_t numSlots, value::TypeTags arrTag, value::Value arrVal);
 
     /**
      * Make a mock scan from an BSON array. This method does NOT assume ownership of the BSONArray
@@ -165,7 +165,7 @@ public:
      * ownership of the BSONArray passed in.
      */
     std::pair<value::SlotVector, std::unique_ptr<PlanStage>> generateMockScanMulti(
-        int64_t numSlots, const BSONArray& array);
+        int32_t numSlots, const BSONArray& array);
 
     /**
      * Prepares the tree of PlanStages given by `root` and returns the SlotAccessor* for `slot`.
@@ -232,7 +232,7 @@ public:
      * slots. `output` should be an array of subarrays with each subarray having M elements, where M
      * is the number of output slots.
      */
-    void runTestMulti(int64_t numInputSlots,
+    void runTestMulti(int32_t numInputSlots,
                       value::TypeTags inputTag,
                       value::Value inputVal,
                       value::TypeTags expectedTag,
diff --git a/src/mongo/db/exec/sbe/values/value.h b/src/mongo/db/exec/sbe/values/value.h
index 14782dfa3ba..757503f8bf8 100644
--- a/src/mongo/db/exec/sbe/values/value.h
+++ b/src/mongo/db/exec/sbe/values/value.h
@@ -205,6 +205,8 @@ inline std::size_t hashCombine(std::size_t state, std::size_t val) noexcept {
  */
 class ValueGuard {
 public:
+    ValueGuard(const std::pair<TypeTags, Value> typedValue)
+        : ValueGuard(typedValue.first, typedValue.second) {}
     ValueGuard(TypeTags tag, Value val) : _tag(tag), _value(val) {}
     ValueGuard() = delete;
     ValueGuard(const ValueGuard&) = delete;
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index 738ccb4ed92..e3f1e38d3df 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -437,51 +437,92 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::getElement(value::Type
                                                                      value::Value arrValue,
                                                                      value::TypeTags idxTag,
                                                                      value::Value idxValue) {
-    if (arrTag != value::TypeTags::Array && arrTag != value::TypeTags::bsonArray) {
+    // We need to ensure that 'size_t' is wide enough to store 32-bit index.
+    static_assert(sizeof(size_t) >= sizeof(int32_t), "size_t must be at least 32-bits");
+
+    if (!value::isArray(arrTag)) {
         return {false, value::TypeTags::Nothing, 0};
     }
 
-    // Bail out if the `idx` parameter isn't a number, or if it can't be converted to a 64-bit
-    // integer, or if it's outside of the range where the `lhsTag` type can represent consecutive
-    // integers precisely.
-    auto [numTag, numVal] = genericNumConvertToPreciseInt64(idxTag, idxValue);
-    if (numTag != value::TypeTags::NumberInt64) {
+    if (idxTag != value::TypeTags::NumberInt32) {
         return {false, value::TypeTags::Nothing, 0};
     }
-    int64_t numInt64 = value::bitcastTo<int64_t>(numVal);
-    // Cast the `idx` parameter to size_t. Bail out if its negative or if it's too big for size_t.
-    if (numInt64 < 0 ||
-        (sizeof(size_t) < sizeof(int64_t) &&
-         numInt64 > static_cast<int64_t>(std::numeric_limits<size_t>::max()))) {
-        return {false, value::TypeTags::Nothing, 0};
+
+    const auto idxInt32 = value::bitcastTo<int32_t>(idxValue);
+    const bool isNegative = idxInt32 < 0;
+
+    size_t idx = 0;
+    if (isNegative) {
+        // Upcast 'idxInt32' to 'int64_t' prevent overflow during the sign change.
+        idx = static_cast<size_t>(-static_cast<int64_t>(idxInt32));
+    } else {
+        idx = static_cast<size_t>(idxInt32);
     }
-    size_t idx = static_cast<size_t>(numInt64);
 
     if (arrTag == value::TypeTags::Array) {
         // If `arr` is an SBE array, use Array::getAt() to retrieve the element at index `idx`.
-        auto [tag, val] = value::getArrayView(arrValue)->getAt(idx);
+        auto arrayView = value::getArrayView(arrValue);
+
+        size_t convertedIdx = idx;
+        if (isNegative) {
+            if (idx > arrayView->size()) {
+                return {false, value::TypeTags::Nothing, 0};
+            }
+            convertedIdx = arrayView->size() - idx;
+        }
+
+        auto [tag, val] = value::getArrayView(arrValue)->getAt(convertedIdx);
         return {false, tag, val};
-    } else if (arrTag == value::TypeTags::bsonArray) {
-        // If `arr` is a BSON array, loop over the elements until we reach the idx-th element.
-        auto be = value::bitcastTo<const char*>(arrValue);
-        auto end = be + ConstDataView(be).read<LittleEndian<uint32_t>>();
-        // Skip document length.
-        be += 4;
-        // The field names of an array are always be 0 thru N-1 in order. Therefore we don't need to
-        // inspect the field names (aside from determining their length so we can skip over them).
-        for (size_t currentIdx = 0; *be != 0; ++currentIdx) {
-            size_t fieldNameLength = strlen(be + 1);
-            if (currentIdx == idx) {
-                auto [tag, val] = bson::convertFrom(true, be, end, fieldNameLength);
-                return {false, tag, val};
+    } else if (arrTag == value::TypeTags::bsonArray || arrTag == value::TypeTags::ArraySet) {
+        value::ArrayEnumerator enumerator(arrTag, arrValue);
+
+        if (!isNegative) {
+            // Loop through array until we meet element at position 'idx'.
+            size_t i = 0;
+            while (i < idx && !enumerator.atEnd()) {
+                i++;
+                enumerator.advance();
+            }
+            // If the array didn't have an element at index `idx`, return Nothing.
+            if (enumerator.atEnd()) {
+                return {false, value::TypeTags::Nothing, 0};
             }
-            be = bson::advance(be, fieldNameLength);
+            auto [tag, val] = enumerator.getViewOfValue();
+            return {false, tag, val};
         }
-        // If the array didn't have an element at index `idx`, return Nothing.
-        return {false, value::TypeTags::Nothing, 0};
+
+        // For negative indexes we use two pointers approach. We start two array enumerators at the
+        // distance of 'idx' and move them at the same time. Once one of the enumerators reaches the
+        // end of the array, the second one points to the element at position '-idx'.
+        //
+        // First, move one of the enumerators 'idx' elements forward.
+        size_t i = 0;
+        while (i < idx && !enumerator.atEnd()) {
+            enumerator.advance();
+            i++;
+        }
+
+        if (i != idx) {
+            // Array is too small to have an element at the requested index.
+            return {false, value::TypeTags::Nothing, 0};
+        }
+
+        // Initiate second enumerator at the start of the array. Now the distance between
+        // 'enumerator' and 'windowEndEnumerator' is exactly 'idx' elements. Move both enumerators
+        // until the first one reaches the end of the array.
+        value::ArrayEnumerator windowEndEnumerator(arrTag, arrValue);
+        while (!enumerator.atEnd() && !windowEndEnumerator.atEnd()) {
+            enumerator.advance();
+            windowEndEnumerator.advance();
+        }
+        invariant(enumerator.atEnd());
+        invariant(!windowEndEnumerator.atEnd());
+
+        auto [tag, val] = windowEndEnumerator.getViewOfValue();
+        return {false, tag, val};
     } else {
-        // Earlier in this function we bailed out if the `arrTag` wasn't Array or bsonArray, so it
-        // should be impossible to reach this point.
+        // Earlier in this function we bailed out if the `arrTag` wasn't Array, ArraySet or
+        // bsonArray, so it should be impossible to reach this point.
         MONGO_UNREACHABLE
     }
 }
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index 5f1f84f04cc..ebd5c607541 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -475,6 +475,38 @@ std::unique_ptr<sbe::EExpression> generateNullishOrNotRepresentableInt32Check(
             sbe::makeE<sbe::EFunction>("exists", sbe::makeEs(std::move(numericConvert32)))));
 }
 
+std::unique_ptr<sbe::EExpression> makeNot(std::unique_ptr<sbe::EExpression> e) {
+    return sbe::makeE<sbe::EPrimUnary>(sbe::EPrimUnary::logicNot, std::move(e));
+}
+
+void buildArrayAccessByConstantIndex(ExpressionVisitorContext* context,
+                                     const std::string& exprName,
+                                     int32_t index) {
+    context->ensureArity(1);
+
+    auto array = context->popExpr();
+
+    auto frameId = context->frameIdGenerator->generate();
+    auto binds = sbe::makeEs(std::move(array));
+    sbe::EVariable arrayRef{frameId, 0};
+
+    auto indexExpr = sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt32,
+                                                sbe::value::bitcastFrom<int32_t>(index));
+    auto argumentIsNotArray =
+        makeNot(sbe::makeE<sbe::EFunction>("isArray", sbe::makeEs(arrayRef.clone())));
+    auto resultExpr = buildMultiBranchConditional(
+        CaseValuePair{generateNullOrMissing(arrayRef),
+                      sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Null, 0)},
+        CaseValuePair{std::move(argumentIsNotArray),
+                      sbe::makeE<sbe::EFail>(ErrorCodes::Error{5126704},
+                                             exprName + " argument must be an array")},
+        sbe::makeE<sbe::EFunction>("getElement",
+                                   sbe::makeEs(arrayRef.clone(), std::move(indexExpr))));
+
+    context->pushExpr(
+        sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(resultExpr)));
+}
+
 class ExpressionPreVisitor final : public ExpressionVisitor {
 public:
     ExpressionPreVisitor(ExpressionVisitorContext* context) : _context{context} {}
@@ -1122,13 +1154,60 @@ public:
         unsupportedExpression(expr->getOpName());
     }
     void visit(ExpressionArrayElemAt* expr) final {
-        unsupportedExpression(expr->getOpName());
+        _context->ensureArity(2);
+
+        auto index = _context->popExpr();
+        auto array = _context->popExpr();
+
+        auto frameId = _context->frameIdGenerator->generate();
+        auto binds = sbe::makeEs(std::move(array), std::move(index));
+        sbe::EVariable arrayRef{frameId, 0};
+        sbe::EVariable indexRef{frameId, 1};
+
+        auto int32Index = [&]() {
+            auto convertedIndex = sbe::makeE<sbe::ENumericConvert>(
+                indexRef.clone(), sbe::value::TypeTags::NumberInt32);
+            auto frameId = _context->frameIdGenerator->generate();
+            auto binds = sbe::makeEs(std::move(convertedIndex));
+            sbe::EVariable convertedIndexRef{frameId, 0};
+
+            auto inExpression = sbe::makeE<sbe::EIf>(
+                sbe::makeE<sbe::EFunction>("exists", sbe::makeEs(convertedIndexRef.clone())),
+                convertedIndexRef.clone(),
+                sbe::makeE<sbe::EFail>(
+                    ErrorCodes::Error{5126703},
+                    "$arrayElemAt second argument cannot be represented as a 32-bit integer"));
+
+            return sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(inExpression));
+        }();
+
+        auto anyOfArgumentsIsNullish =
+            sbe::makeE<sbe::EPrimBinary>(sbe::EPrimBinary::logicOr,
+                                         generateNullOrMissing(arrayRef),
+                                         generateNullOrMissing(indexRef));
+        auto firstArgumentIsNotArray =
+            makeNot(sbe::makeE<sbe::EFunction>("isArray", sbe::makeEs(arrayRef.clone())));
+        auto secondArgumentIsNotNumeric = generateNonNumericCheck(indexRef);
+        auto arrayElemAtExpr = buildMultiBranchConditional(
+            CaseValuePair{std::move(anyOfArgumentsIsNullish),
+                          sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::Null, 0)},
+            CaseValuePair{std::move(firstArgumentIsNotArray),
+                          sbe::makeE<sbe::EFail>(ErrorCodes::Error{5126701},
+                                                 "$arrayElemAt first argument must be an array")},
+            CaseValuePair{std::move(secondArgumentIsNotNumeric),
+                          sbe::makeE<sbe::EFail>(ErrorCodes::Error{5126702},
+                                                 "$arrayElemAt second argument must be a number")},
+            sbe::makeE<sbe::EFunction>("getElement",
+                                       sbe::makeEs(arrayRef.clone(), std::move(int32Index))));
+
+        _context->pushExpr(
+            sbe::makeE<sbe::ELocalBind>(frameId, std::move(binds), std::move(arrayElemAtExpr)));
     }
     void visit(ExpressionFirst* expr) final {
-        unsupportedExpression(expr->getOpName());
+        buildArrayAccessByConstantIndex(_context, expr->getOpName(), 0);
     }
     void visit(ExpressionLast* expr) final {
-        unsupportedExpression(expr->getOpName());
+        buildArrayAccessByConstantIndex(_context, expr->getOpName(), -1);
     }
     void visit(ExpressionObjectToArray* expr) final {
         unsupportedExpression(expr->getOpName());