SERVER-6074 Add $slice aggregation expression

jstests/aggregation/bugs/server6074.js

@@ -0,0 +1,81 @@
+// SERVER-6074: Add $slice aggregation expression.
+
+// For assertErrorCode.
+load('jstests/aggregation/extras/utils.js');
+
+(function() {
+    'use strict';
+
+    var coll = db.agg_slice_expr;
+    coll.drop();
+
+    // Need to have at least one document to ensure the pipeline executes.
+    assert.writeOK(coll.insert({}));
+
+    function testSlice(sliceArgs, expArray) {
+        var pipeline = [{$project: {_id: 0, slice: {$slice: sliceArgs}}}];
+        assert.eq(coll.aggregate(pipeline).toArray(), [{slice: expArray}]);
+    }
+
+    // Two argument form.
+
+    testSlice([[0, 1, 2, 3, 4], 2], [0, 1]);
+    testSlice([[0, 1, 2, 3, 4], 2.0], [0, 1]);
+    // Negative count
+    testSlice([[0, 1, 2, 3, 4], -2], [3, 4]);
+    testSlice([[0, 1, 2, 3, 4], -2.0], [3, 4]);
+    // Zero count.
+    testSlice([[0, 1, 2, 3, 4], 0], []);
+    // Out of bounds positive.
+    testSlice([[0, 1, 2, 3, 4], 10], [0, 1, 2, 3, 4]);
+    // Out of bounds negative.
+    testSlice([[0, 1, 2, 3, 4], -10], [0, 1, 2, 3, 4]);
+    // Null arguments
+    testSlice([null, -10], null);
+    testSlice([[0, 1, 2, 3, 4], null], null);
+
+    // Three argument form.
+
+    testSlice([[0, 1, 2, 3, 4], 1, 2], [1, 2]);
+    testSlice([[0, 1, 2, 3, 4], 1.0, 2.0], [1, 2]);
+    // Negative start index.
+    testSlice([[0, 1, 2, 3, 4], -3, 2], [2, 3]);
+    testSlice([[0, 1, 2, 3, 4], -5, 2], [0, 1]);
+    // Slice starts out of bounds.
+    testSlice([[0, 1, 2, 3, 4], -10, 2], [0, 1]);
+    testSlice([[0, 1, 2, 3, 4], 10, 2], []);
+    // Slice ends out of bounds.
+    testSlice([[0, 1, 2, 3, 4], 4, 3], [4]);
+    testSlice([[0, 1, 2, 3, 4], -1, 3], [4]);
+    // Null arguments
+    testSlice([[0, 1, 2, 3, 4], -1, null], null);
+
+    // Error cases.
+
+    // Wrong number of arguments.
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2, 3]]}}}], 28667);
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2, 3], 4, 5, 6]}}}], 28667);
+
+    // First argument is not an array.
+    assertErrorCode(coll, [{$project: {x: {$slice: ['one', 2]}}}], 28724);
+
+    // Second argument is not numeric.
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], '2']}}}], 28725);
+
+    // Second argument is not integral.
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], 1.5]}}}], 28726);
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], Math.pow(2, 32)]}}}], 28726);
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], -Math.pow(2, 31) - 1]}}}], 28726);
+
+    // Third argument is not numeric.
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], 0, '2']}}}], 28727);
+
+    // Third argument is not integral.
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], 0, 1.5]}}}], 28728);
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], 0, Math.pow(2, 32)]}}}], 28728);
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], 0, -Math.pow(2, 31) - 1]}}}], 28728);
+
+    // Third argument is not positive.
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], 0, 0]}}}], 28729);
+    assertErrorCode(coll, [{$project: {x: {$slice: [[1, 2], 0, -1]}}}], 28729);
+}());

src/mongo/db/pipeline/expression.cpp

@@ -30,8 +30,10 @@
 
 #include "mongo/db/pipeline/expression.h"
 
+#include <algorithm>
 #include <boost/algorithm/string.hpp>
 #include <cstdio>
+#include <vector>
 
 #include "mongo/db/jsobj.h"
 #include "mongo/db/pipeline/document.h"
@@ -2433,6 +2435,92 @@ const char* ExpressionIsArray::getOpName() const {
     return "$isArray";
 }
 
+/* ----------------------- ExpressionSlice ---------------------------- */
+
+Value ExpressionSlice::evaluateInternal(Variables* vars) const {
+    const size_t n = vpOperand.size();
+
+    Value arrayVal = vpOperand[0]->evaluateInternal(vars);
+    // Could be either a start index or the length from 0.
+    Value arg2 = vpOperand[1]->evaluateInternal(vars);
+
+    if (arrayVal.nullish() || arg2.nullish()) {
+        return Value(BSONNULL);
+    }
+
+    uassert(28724,
+            str::stream() << "First argument to $slice must be an Array, but is"
+                          << " of type: " << typeName(arrayVal.getType()),
+            arrayVal.getType() == Array);
+    uassert(28725,
+            str::stream() << "Second argument to $slice must be a numeric value,"
+                          << " but is of type: " << typeName(arg2.getType()),
+            arg2.numeric());
+    uassert(28726,
+            str::stream() << "Second argument to $slice can't be represented as"
+                          << " a 32-bit integer: " << arg2.coerceToDouble(),
+            arg2.integral());
+
+    const auto& array = arrayVal.getArray();
+    size_t start;
+    size_t end;
+
+    if (n == 2) {
+        // Only count given.
+        int count = arg2.coerceToInt();
+        start = 0;
+        end = array.size();
+        if (count >= 0) {
+            end = std::min(end, size_t(count));
+        } else {
+            // Negative count's start from the back. If a abs(count) is greater
+            // than the
+            // length of the array, return the whole array.
+            start = std::max(0, static_cast<int>(array.size()) + count);
+        }
+    } else {
+        // We have both a start index and a count.
+        int startInt = arg2.coerceToInt();
+        if (startInt < 0) {
+            // Negative values start from the back. If a abs(start) is greater
+            // than the length
+            // of the array, start from 0.
+            start = std::max(0, static_cast<int>(array.size()) + startInt);
+        } else {
+            start = std::min(array.size(), size_t(startInt));
+        }
+
+        Value countVal = vpOperand[2]->evaluateInternal(vars);
+
+        if (countVal.nullish()) {
+            return Value(BSONNULL);
+        }
+
+        uassert(28727,
+                str::stream() << "Third argument to $slice must be numeric, but "
+                              << "is of type: " << typeName(countVal.getType()),
+                countVal.numeric());
+        uassert(28728,
+                str::stream() << "Third argument to $slice can't be represented"
+                              << " as a 32-bit integer: " << countVal.coerceToDouble(),
+                countVal.integral());
+        uassert(28729,
+                str::stream() << "Third argument to $slice must be positive: "
+                              << countVal.coerceToInt(),
+                countVal.coerceToInt() > 0);
+
+        size_t count = size_t(countVal.coerceToInt());
+        end = std::min(start + count, array.size());
+    }
+
+    return Value(vector<Value>(array.begin() + start, array.begin() + end));
+}
+
+REGISTER_EXPRESSION(slice, ExpressionSlice::parse);
+const char* ExpressionSlice::getOpName() const {
+    return "$slice";
+}
+
 /* ----------------------- ExpressionSize ---------------------------- */
 
 Value ExpressionSize::evaluateInternal(Variables* vars) const {

src/mongo/db/pipeline/expression.h

@@ -382,6 +382,23 @@ public:
 template <typename SubClass>
 class ExpressionVariadic : public ExpressionNaryBase<SubClass> {};
 
+/**
+ * Inherit from this class if your expression can take a range of arguments, e.g. if it has some
+ * optional arguments.
+ */
+template <typename SubClass, int MinArgs, int MaxArgs>
+class ExpressionRangedArity : public ExpressionNaryBase<SubClass> {
+public:
+    void validateArguments(const Expression::ExpressionVector& args) const override {
+        uassert(28667,
+                mongoutils::str::stream()
+                    << "Expression " << this->getOpName() << " takes at least " << MinArgs
+                    << " arguments, and at most " << MaxArgs << ", but " << args.size()
+                    << " were passed in.",
+                MinArgs <= args.size() && args.size() <= MaxArgs);
+    }
+};
+
 /// Inherit from this class if your expression takes a fixed number of arguments.
 template <typename SubClass, int NArgs>
 class ExpressionFixedArity : public ExpressionNaryBase<SubClass> {
@@ -1001,6 +1018,20 @@ public:
 };
 
 
+class ExpressionSize final : public ExpressionFixedArity<ExpressionSize, 1> {
+public:
+    Value evaluateInternal(Variables* vars) const final;
+    const char* getOpName() const final;
+};
+
+
+class ExpressionSlice final : public ExpressionRangedArity<ExpressionSlice, 2, 3> {
+public:
+    Value evaluateInternal(Variables* vars) const final;
+    const char* getOpName() const final;
+};
+
+
 class ExpressionIsArray final : public ExpressionFixedArity<ExpressionIsArray, 1> {
 public:
     Value evaluateInternal(Variables* vars) const final;
@@ -1008,17 +1039,12 @@ public:
 };
 
 
-class ExpressionSize final : public ExpressionFixedArity<ExpressionSize, 1> {
-public:
-    Value evaluateInternal(Variables* vars) const final;
-    const char* getOpName() const final;
-};
-
 class ExpressionSqrt final : public ExpressionFixedArity<ExpressionSqrt, 1> {
     Value evaluateInternal(Variables* vars) const final;
     const char* getOpName() const final;
 };
 
+
 class ExpressionStrcasecmp final : public ExpressionFixedArity<ExpressionStrcasecmp, 2> {
 public:
     Value evaluateInternal(Variables* vars) const final;