SERVER-51531: Implement $array expr

8 files changed, 78 insertions, 5 deletions

diff --git a/jstests/aggregation/expressions/array_expression.js b/jstests/aggregation/expressions/array_expression.js
new file mode 100644
index 00000000000..8eb9c06c0e2
--- /dev/null
+++ b/jstests/aggregation/expressions/array_expression.js
@@ -0,0 +1,33 @@
+// Tests for $array expression.
+(function() {
+"use strict";
+
+let coll = db.array_expr;
+coll.drop();
+
+function assertArray(expArray, ...inputs) {
+    assert(coll.drop());
+    if (inputs.length == 0) {
+        assert.commandWorked(coll.insert({}));
+    } else if (inputs.length == 1) {
+        assert.commandWorked(coll.insert({a: inputs[0]}));
+    } else {
+        assert.commandWorked(coll.insert({a: inputs[0], b: inputs[1]}));
+    }
+    const result = coll.aggregate([{$project: {out: ["$a", "$b"]}}]).toArray()[0].out;
+    assert.eq(result, expArray);
+}
+
+assertArray([1, 2], 1, 2);
+assertArray([null, null], null, null);
+assertArray(["TestInput", null], "TestInput", null);
+assertArray([{a: 1, b: 2}, [1, 2]], {a: 1, b: 2}, [1, 2]);
+assertArray(["TestInput", null], "TestInput");
+assertArray([null, null]);
+
+// no arg
+assert(coll.drop());
+assert.commandWorked(coll.insert({}));
+let result = coll.aggregate([{$project: {out: []}}]).toArray()[0].out;
+assert.eq(result, []);
+}());
diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index a5693e47668..d3cb0c35cea 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -779,6 +779,7 @@ static stdx::unordered_map<std::string, BuiltinFn> kBuiltinFunctions = {
      BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::objectToArray, false}},
     {"arrayToObject",
      BuiltinFn{[](size_t n) { return n == 1; }, vm::Builtin::arrayToObject, false}},
+    {"array", BuiltinFn{kAnyNumberOfArgs, vm::Builtin::newArray, false}},
 };
 
 /**
diff --git a/src/mongo/db/exec/sbe/vm/vm.cpp b/src/mongo/db/exec/sbe/vm/vm.cpp
index bbf03060e3a..7c8e8171d9d 100644
--- a/src/mongo/db/exec/sbe/vm/vm.cpp
+++ b/src/mongo/db/exec/sbe/vm/vm.cpp
@@ -2027,9 +2027,8 @@ FastTuple<bool, value::TypeTags, value::Value> ByteCode::builtinNewArray(ArityTy
     if (arity) {
         arr->reserve(arity);
         for (ArityType idx = 0; idx < arity; ++idx) {
-            auto [owned, tag, val] = getFromStack(idx);
-            auto [tagCopy, valCopy] = value::copyValue(tag, val);
-            arr->push_back(tagCopy, valCopy);
+            auto [tag, val] = moveOwnedFromStack(idx);
+            arr->push_back(tag, val);
         }
     }
 
diff --git a/src/mongo/db/pipeline/accumulator_multi.cpp b/src/mongo/db/pipeline/accumulator_multi.cpp
index a8dedfcc086..9d2ed9ae259 100644
--- a/src/mongo/db/pipeline/accumulator_multi.cpp
+++ b/src/mongo/db/pipeline/accumulator_multi.cpp
@@ -510,6 +510,7 @@ std::pair<SortPattern, BSONArray> parseAccumulatorTopBottomNSortBy(ExpressionCon
 template <TopBottomSense sense, bool single>
 AccumulationExpression AccumulatorTopBottomN<sense, single>::parseTopBottomN(
     ExpressionContext* const expCtx, BSONElement elem, VariablesParseState vps) {
+    expCtx->sbeGroupCompatibility = SbeCompatibility::notCompatible;
     auto name = AccumulatorTopBottomN<sense, single>::getName();
     const auto [n, output, sortBy] =
         accumulatorNParseArgs<single>(expCtx, elem, name.rawData(), true, vps);
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 91341d79546..431187dab68 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1128,7 +1128,8 @@ class ExpressionArray final : public ExpressionVariadic<ExpressionArray> {
 public:
     explicit ExpressionArray(ExpressionContext* const expCtx)
         : ExpressionVariadic<ExpressionArray>(expCtx) {
-        expCtx->sbeCompatibility = SbeCompatibility::notCompatible;
+        expCtx->sbeCompatibility =
+            std::min(expCtx->sbeCompatibility, SbeCompatibility::flagGuarded);
     }
 
     ExpressionArray(ExpressionContext* const expCtx,
diff --git a/src/mongo/db/query/sbe_stage_builder_abt_helpers.cpp b/src/mongo/db/query/sbe_stage_builder_abt_helpers.cpp
index f90e1b1b5e0..8b023cb89c5 100644
--- a/src/mongo/db/query/sbe_stage_builder_abt_helpers.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_abt_helpers.cpp
@@ -140,6 +140,11 @@ optimizer::ABT makeFillEmptyTrue(optimizer::ABT e) {
     return makeFillEmpty(std::move(e), true);
 }
 
+optimizer::ABT makeFillEmptyNull(optimizer::ABT e) {
+    return optimizer::make<optimizer::BinaryOp>(
+        optimizer::Operations::FillEmpty, std::move(e), optimizer::Constant::null());
+}
+
 optimizer::ABT makeNot(optimizer::ABT e) {
     return makeUnaryOp(optimizer::Operations::Not, std::move(e));
 }
diff --git a/src/mongo/db/query/sbe_stage_builder_abt_helpers.h b/src/mongo/db/query/sbe_stage_builder_abt_helpers.h
index c3099c16aec..4b7ef4eb289 100644
--- a/src/mongo/db/query/sbe_stage_builder_abt_helpers.h
+++ b/src/mongo/db/query/sbe_stage_builder_abt_helpers.h
@@ -75,6 +75,10 @@ optimizer::ABT makeFillEmptyFalse(optimizer::ABT e);
  * expression.
  */
 optimizer::ABT makeFillEmptyTrue(optimizer::ABT e);
+/**
+ * Check if expression returns Nothing and return null if so. Otherwise, return the expression.
+ */
+optimizer::ABT makeFillEmptyNull(optimizer::ABT e);
 optimizer::ABT makeNot(optimizer::ABT e);
 
 optimizer::ProjectionName makeVariableName(sbe::value::SlotId slotId);
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index 98e26148812..8415939b533 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -820,7 +820,36 @@ public:
     }
 
     void visit(const ExpressionArray* expr) final {
-        unsupportedExpression(expr->getOpName());
+        auto arity = expr->getChildren().size();
+        _context->ensureArity(arity);
+
+        if (arity == 0) {
+            auto [emptyArrTag, emptyArrVal] = sbe::value::makeNewArray();
+            pushABT(makeABTConstant(emptyArrTag, emptyArrVal));
+            return;
+        }
+
+        std::vector<std::pair<optimizer::ProjectionName, optimizer::ABT>> binds;
+        for (size_t idx = 0; idx < arity; ++idx) {
+            binds.emplace_back(makeLocalVariableName(_context->state.frameId(), 0),
+                               _context->popABTExpr());
+        }
+        std::reverse(std::begin(binds), std::end(binds));
+
+        optimizer::ABTVector argVars;
+        for (auto& bind : binds) {
+            argVars.push_back(makeFillEmptyNull(makeVariable(bind.first)));
+        }
+
+        auto arrayExpr = optimizer::make<optimizer::FunctionCall>("newArray", std::move(argVars));
+
+        for (auto it = binds.begin(); it != binds.end(); it++) {
+            arrayExpr = optimizer::make<optimizer::Let>(
+                it->first, std::move(it->second), std::move(arrayExpr));
+        }
+
+        pushABT(std::move(arrayExpr));
+        return;
     }
     void visit(const ExpressionArrayElemAt* expr) final {
         unsupportedExpression(expr->getOpName());