SERVER-4589: Add $arrayElemAt aggregation expression

4 files changed, 70 insertions, 5 deletions

diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 8c266d22158..34cbda702a3 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -593,6 +593,44 @@ namespace {
         return "$anyElementTrue";
     }
 
+    /* ------------------------- ExpressionArrayElemAt -------------------------- */
+
+    Value ExpressionArrayElemAt::evaluateInternal(Variables* vars) const {
+        const Value array = vpOperand[0]->evaluateInternal(vars);
+        const Value indexArg = vpOperand[1]->evaluateInternal(vars);
+
+        if (array.nullish() || indexArg.nullish()) {
+            return Value(BSONNULL);
+        }
+
+        uassert(28689, str::stream() << getOpName() << "'s first argument must be an array, but is "
+                                     << typeName(array.getType()),
+                array.getType() == Array);
+        uassert(28690, str::stream() << 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"
+                                     << " a 32-bit integer: " << indexArg.coerceToDouble(),
+                indexArg.integral());
+
+        long long i = indexArg.coerceToLong();
+        if (i < 0 && static_cast<size_t>(std::abs(i)) > array.getArrayLength()) {
+            // Positive indices that are too large are handled automatically by Value.
+            return Value();
+        }
+        else if (i < 0) {
+            // Index from the back of the array.
+            i = array.getArrayLength() + i;
+        }
+        const size_t index = static_cast<size_t>(i);
+        return array[index];
+    }
+
+    REGISTER_EXPRESSION("$arrayElemAt", ExpressionArrayElemAt::parse);
+    const char* ExpressionArrayElemAt::getOpName() const {
+        return "$arrayElemAt";
+    }
+
     /* -------------------- ExpressionCoerceToBool ------------------------- */
 
     intrusive_ptr<ExpressionCoerceToBool> ExpressionCoerceToBool::create(
@@ -1931,11 +1969,8 @@ namespace {
             uassert(16610, "can't $mod by 0",
                     right != 0);
 
-            if (leftType == NumberDouble
-                || (rightType == NumberDouble && rhs.coerceToInt() != right)) {
-                // the shell converts ints to doubles so if right is larger than int max or
-                // if right truncates to something other than itself, it is a real double.
-                // Integer-valued double case is handled below
+            if (leftType == NumberDouble || (rightType == NumberDouble && !rhs.integral())) {
+                // Need to do fmod. Integer-valued double case is handled below.
 
                 double left = lhs.coerceToDouble();
                 return Value(fmod(left, right));
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index dfdadf0cccb..258e0d6c034 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -403,6 +403,13 @@ namespace mongo {
     };
 
 
+    class ExpressionArrayElemAt final : public ExpressionFixedArity<ExpressionArrayElemAt, 2> {
+    public:
+        Value evaluateInternal(Variables* vars) const final;
+        const char* getOpName() const final;
+    };
+
+
     class ExpressionCoerceToBool : public Expression {
     public:
         // virtuals from ExpressionNary
diff --git a/src/mongo/db/pipeline/value.cpp b/src/mongo/db/pipeline/value.cpp
index 833d7ac23de..0d55a309002 100644
--- a/src/mongo/db/pipeline/value.cpp
+++ b/src/mongo/db/pipeline/value.cpp
@@ -31,6 +31,7 @@
 #include "mongo/db/pipeline/value.h"
 
 #include <cmath>
+#include <limits>
 #include <boost/functional/hash.hpp>
 
 #include "mongo/base/compare_numbers.h"
@@ -822,6 +823,22 @@ namespace mongo {
         return Undefined;
     }
 
+    bool Value::integral() const {
+        switch (getType()) {
+        case NumberInt:
+            return true;
+        case NumberLong:
+            return (_storage.longValue <= numeric_limits<int>::max()
+                    && _storage.longValue >= numeric_limits<int>::min());
+        case NumberDouble:
+            return (_storage.doubleValue <= numeric_limits<int>::max()
+                    && _storage.doubleValue >= numeric_limits<int>::min()
+                    && _storage.doubleValue == static_cast<int>(_storage.doubleValue));
+        default:
+            return false;
+        }
+    }
+
     size_t Value::getApproximateSize() const {
         switch(getType()) {
         case Code:
diff --git a/src/mongo/db/pipeline/value.h b/src/mongo/db/pipeline/value.h
index f388b82be7b..41201f64873 100644
--- a/src/mongo/db/pipeline/value.h
+++ b/src/mongo/db/pipeline/value.h
@@ -124,6 +124,12 @@ namespace mongo {
                 || _storage.type == NumberInt;
         }
 
+        /**
+         * Returns true if this value is a numeric type that can be represented as a 32-bit integer,
+         * and false otherwise.
+         */
+        bool integral() const;
+
         /// Get the BSON type of the field.
         BSONType getType() const { return _storage.bsonType(); }