SERVER-44327 Add $first and $last expressions

4 files changed, 151 insertions, 12 deletions

diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index ee960534f0e..3c557d7d768 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -510,24 +510,24 @@ const char* ExpressionArray::getOpName() const {
 
 /* ------------------------- ExpressionArrayElemAt -------------------------- */
 
-Value ExpressionArrayElemAt::evaluate(const Document& root, Variables* variables) const {
-    const Value array = _children[0]->evaluate(root, variables);
-    const Value indexArg = _children[1]->evaluate(root, variables);
-
+namespace {
+Value arrayElemAt(const ExpressionNary* self, Value array, Value indexArg) {
     if (array.nullish() || indexArg.nullish()) {
         return Value(BSONNULL);
     }
 
+    size_t arity = self->getOperandList().size();
     uassert(28689,
-            str::stream() << getOpName() << "'s first argument must be an array, but is "
-                          << typeName(array.getType()),
+            str::stream() << self->getOpName() << "'s "
+                          << (arity == 1 ? "argument" : "first argument")
+                          << " must be an array, but is " << typeName(array.getType()),
             array.isArray());
     uassert(28690,
-            str::stream() << getOpName() << "'s second argument must be a numeric value,"
+            str::stream() << self->getOpName() << "'s second argument must be a numeric value,"
                           << " but is " << typeName(indexArg.getType()),
             indexArg.numeric());
     uassert(28691,
-            str::stream() << getOpName() << "'s second argument must be representable as"
+            str::stream() << self->getOpName() << "'s second argument must be representable as"
                           << " a 32-bit integer: " << indexArg.coerceToDouble(),
             indexArg.integral());
 
@@ -542,12 +542,43 @@ Value ExpressionArrayElemAt::evaluate(const Document& root, Variables* variables
     const size_t index = static_cast<size_t>(i);
     return array[index];
 }
+}  // namespace
+
+Value ExpressionArrayElemAt::evaluate(const Document& root, Variables* variables) const {
+    const Value array = _children[0]->evaluate(root, variables);
+    const Value indexArg = _children[1]->evaluate(root, variables);
+    return arrayElemAt(this, array, indexArg);
+}
 
 REGISTER_EXPRESSION(arrayElemAt, ExpressionArrayElemAt::parse);
 const char* ExpressionArrayElemAt::getOpName() const {
     return "$arrayElemAt";
 }
 
+/* ------------------------- ExpressionFirst -------------------------- */
+
+Value ExpressionFirst::evaluate(const Document& root, Variables* variables) const {
+    const Value array = _children[0]->evaluate(root, variables);
+    return arrayElemAt(this, array, Value(0));
+}
+
+REGISTER_EXPRESSION(first, ExpressionFirst::parse);
+const char* ExpressionFirst::getOpName() const {
+    return "$first";
+}
+
+/* ------------------------- ExpressionLast -------------------------- */
+
+Value ExpressionLast::evaluate(const Document& root, Variables* variables) const {
+    const Value array = _children[0]->evaluate(root, variables);
+    return arrayElemAt(this, array, Value(-1));
+}
+
+REGISTER_EXPRESSION(last, ExpressionLast::parse);
+const char* ExpressionLast::getOpName() const {
+    return "$last";
+}
+
 /* ------------------------- ExpressionObjectToArray -------------------------- */
 
 Value ExpressionObjectToArray::evaluate(const Document& root, Variables* variables) const {
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index c2dae92cffb..c1957100264 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -864,6 +864,32 @@ public:
     }
 };
 
+class ExpressionFirst final : public ExpressionFixedArity<ExpressionFirst, 1> {
+public:
+    explicit ExpressionFirst(const boost::intrusive_ptr<ExpressionContext>& expCtx)
+        : ExpressionFixedArity<ExpressionFirst, 1>(expCtx) {}
+
+    Value evaluate(const Document& root, Variables* variables) const final;
+    const char* getOpName() const final;
+
+    void acceptVisitor(ExpressionVisitor* visitor) final {
+        return visitor->visit(this);
+    }
+};
+
+class ExpressionLast final : public ExpressionFixedArity<ExpressionLast, 1> {
+public:
+    explicit ExpressionLast(const boost::intrusive_ptr<ExpressionContext>& expCtx)
+        : ExpressionFixedArity<ExpressionLast, 1>(expCtx) {}
+
+    Value evaluate(const Document& root, Variables* variables) const final;
+    const char* getOpName() const final;
+
+    void acceptVisitor(ExpressionVisitor* visitor) final {
+        return visitor->visit(this);
+    }
+};
+
 class ExpressionObjectToArray final : public ExpressionFixedArity<ExpressionObjectToArray, 1> {
 public:
     explicit ExpressionObjectToArray(const boost::intrusive_ptr<ExpressionContext>& expCtx)
diff --git a/src/mongo/db/pipeline/expression_nary_test.cpp b/src/mongo/db/pipeline/expression_nary_test.cpp
index b96220378ce..6fa2e6ec07e 100644
--- a/src/mongo/db/pipeline/expression_nary_test.cpp
+++ b/src/mongo/db/pipeline/expression_nary_test.cpp
@@ -130,11 +130,14 @@ public:
 
 class ExpressionNaryTestOneArg : public ExpressionBaseTest {
 public:
-    virtual void assertEvaluates(Value input, Value output) {
+    Value eval(Value input) {
         addOperand(_expr, input);
-        ASSERT_VALUE_EQ(output, _expr->evaluate({}, &_expr->getExpressionContext()->variables));
-        ASSERT_EQUALS(output.getType(),
-                      _expr->evaluate({}, &_expr->getExpressionContext()->variables).getType());
+        return _expr->evaluate({}, &_expr->getExpressionContext()->variables);
+    }
+    virtual void assertEvaluates(Value input, Value output) {
+        Value v = eval(input);
+        ASSERT_VALUE_EQ(output, v);
+        ASSERT_EQUALS(output.getType(), v.getType());
     }
 
     intrusive_ptr<ExpressionNary> _expr;
@@ -1216,6 +1219,81 @@ TEST_F(ExpressionBinarySizeTest, HandlesBinData) {
 TEST_F(ExpressionBinarySizeTest, HandlesNullish) {
     assertEval(BSONNULL, BSONNULL);
     assertEval(BSONUndefined, BSONNULL);
+    assertEval(Value(), BSONNULL);
+}
+
+/* ------------------------- ExpressionFirst -------------------------- */
+
+class ExpressionFirstTest : public ExpressionNaryTestOneArg {
+public:
+    void assertEval(ImplicitValue input, ImplicitValue output) {
+        intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+        _expr = new ExpressionFirst(expCtx);
+        ExpressionNaryTestOneArg::assertEvaluates(input, output);
+    }
+    void assertEvalFails(ImplicitValue input) {
+        intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+        _expr = new ExpressionFirst(expCtx);
+        ASSERT_THROWS_CODE(eval(input), DBException, 28689);
+    }
+};
+
+TEST_F(ExpressionFirstTest, HandlesArrays) {
+    assertEval(vector<Value>{Value("A"_sd)}, "A"_sd);
+    assertEval(vector<Value>{Value("A"_sd), Value("B"_sd)}, "A"_sd);
+    assertEval(vector<Value>{Value("A"_sd), Value("B"_sd), Value("C"_sd)}, "A"_sd);
+}
+
+TEST_F(ExpressionFirstTest, HandlesEmptyArray) {
+    assertEval(vector<Value>{}, Value());
+}
+
+TEST_F(ExpressionFirstTest, HandlesNullish) {
+    assertEval(BSONNULL, BSONNULL);
+    assertEval(BSONUndefined, BSONNULL);
+    assertEval(Value(), BSONNULL);
+}
+
+TEST_F(ExpressionFirstTest, RejectsNonArrays) {
+    assertEvalFails("asdf"_sd);
+    assertEvalFails(BSONBinData("asdf", 4, BinDataGeneral));
+}
+
+/* ------------------------- ExpressionLast -------------------------- */
+
+class ExpressionLastTest : public ExpressionNaryTestOneArg {
+public:
+    void assertEval(ImplicitValue input, ImplicitValue output) {
+        intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+        _expr = new ExpressionLast(expCtx);
+        ExpressionNaryTestOneArg::assertEvaluates(input, output);
+    }
+    void assertEvalFails(ImplicitValue input) {
+        intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+        _expr = new ExpressionLast(expCtx);
+        ASSERT_THROWS_CODE(eval(input), DBException, 28689);
+    }
+};
+
+TEST_F(ExpressionLastTest, HandlesArrays) {
+    assertEval(vector<Value>{Value("A"_sd)}, "A"_sd);
+    assertEval(vector<Value>{Value("A"_sd), Value("B"_sd)}, "B"_sd);
+    assertEval(vector<Value>{Value("A"_sd), Value("B"_sd), Value("C"_sd)}, "C"_sd);
+}
+
+TEST_F(ExpressionLastTest, HandlesEmptyArray) {
+    assertEval(vector<Value>{}, Value());
+}
+
+TEST_F(ExpressionLastTest, HandlesNullish) {
+    assertEval(BSONNULL, BSONNULL);
+    assertEval(BSONUndefined, BSONNULL);
+    assertEval(Value(), BSONNULL);
+}
+
+TEST_F(ExpressionLastTest, RejectsNonArrays) {
+    assertEvalFails("asdf"_sd);
+    assertEvalFails(BSONBinData("asdf", 4, BinDataGeneral));
 }
 
 }  // anonymous namespace
diff --git a/src/mongo/db/pipeline/expression_visitor.h b/src/mongo/db/pipeline/expression_visitor.h
index 8269da1e1f4..c708d6a4192 100644
--- a/src/mongo/db/pipeline/expression_visitor.h
+++ b/src/mongo/db/pipeline/expression_visitor.h
@@ -45,6 +45,8 @@ class ExpressionAnd;
 class ExpressionAnyElementTrue;
 class ExpressionArray;
 class ExpressionArrayElemAt;
+class ExpressionFirst;
+class ExpressionLast;
 class ExpressionObjectToArray;
 class ExpressionArrayToObject;
 class ExpressionCeil;
@@ -179,6 +181,8 @@ public:
     virtual void visit(ExpressionAnyElementTrue*) = 0;
     virtual void visit(ExpressionArray*) = 0;
     virtual void visit(ExpressionArrayElemAt*) = 0;
+    virtual void visit(ExpressionFirst*) = 0;
+    virtual void visit(ExpressionLast*) = 0;
     virtual void visit(ExpressionObjectToArray*) = 0;
     virtual void visit(ExpressionArrayToObject*) = 0;
     virtual void visit(ExpressionCeil*) = 0;