SERVER-20169 Aggregation now supports the range expressin.

4 files changed, 161 insertions, 0 deletions

diff --git a/jstests/aggregation/bugs/server20169.js b/jstests/aggregation/bugs/server20169.js
new file mode 100644
index 00000000000..f4d8735701e
--- /dev/null
+++ b/jstests/aggregation/bugs/server20169.js
@@ -0,0 +1,61 @@
+// In SERVER-20169, the $range expression was added to the aggregation framework. In this file, we
+// test the behavior and error cases of this expression.
+load("jstests/aggregation/extras/utils.js");  // For assertErrorCode.
+
+(function() {
+    "use strict";
+
+    var coll = db.range;
+    coll.drop();
+
+    // We need an input document to receive an output document.
+    coll.insert({});
+
+    var rangeObj = [1];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    28667,
+                    "range requires two" + " or three arguments");
+
+    rangeObj = ["a", 1];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    34443,
+                    "range requires a" + " numeric starting value");
+
+    rangeObj = [1.1, 1];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    34444,
+                    "range requires an" + " integral starting value");
+
+    rangeObj = [1, "a"];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    34445,
+                    "range requires a" + " numeric ending value");
+
+    rangeObj = [1, 1.1];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    34446,
+                    "range requires an" + " integral ending value");
+
+    rangeObj = [1, 3, "a"];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    34447,
+                    "range requires a" + " numeric step value");
+
+    rangeObj = [1, 3, 1.1];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    34448,
+                    "range requires an" + " integral step value");
+
+    rangeObj = [1, 3, 0];
+    assertErrorCode(coll,
+                        [{$project: {range: {$range: rangeObj}}}],
+                    34449,
+                    "range requires a" + " non-zero step value");
+}());
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 7cdb9a30fab..c14ce934a43 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -2499,6 +2499,65 @@ const char* ExpressionPow::getOpName() const {
     return "$pow";
 }
 
+/* ------------------------- ExpressionRange ------------------------------ */
+
+Value ExpressionRange::evaluateInternal(Variables* vars) const {
+    Value startVal(vpOperand[0]->evaluateInternal(vars));
+    Value endVal(vpOperand[1]->evaluateInternal(vars));
+
+    uassert(34443,
+            str::stream() << "$range requires a numeric starting value, found value of type: "
+                          << typeName(startVal.getType()),
+            startVal.numeric());
+    uassert(34444,
+            str::stream() << "$range requires a starting value that can be represented as a 32-bit "
+                             "integer, found value: " << startVal.toString(),
+            startVal.integral());
+    uassert(34445,
+            str::stream() << "$range requires a numeric ending value, found value of type: "
+                          << typeName(endVal.getType()),
+            endVal.numeric());
+    uassert(34446,
+            str::stream() << "$range requires an ending value that can be represented as a 32-bit "
+                             "integer, found value: " << endVal.toString(),
+            endVal.integral());
+
+    int current = startVal.coerceToInt();
+    int end = endVal.coerceToInt();
+
+    int step = 1;
+    if (vpOperand.size() == 3) {
+        // A step was specified by the user.
+        Value stepVal(vpOperand[2]->evaluateInternal(vars));
+
+        uassert(34447,
+                str::stream() << "$range requires a numeric starting value, found value of type:"
+                              << typeName(stepVal.getType()),
+                stepVal.numeric());
+        uassert(34448,
+                str::stream() << "$range requires a step value that can be represented as a 32-bit "
+                                 "integer, found value: " << stepVal.toString(),
+                stepVal.integral());
+        step = stepVal.coerceToInt();
+
+        uassert(34449, "$range requires a non-zero step value", step != 0);
+    }
+
+    std::vector<Value> output;
+
+    while ((step > 0 ? current < end : current > end)) {
+        output.push_back(Value(current));
+        current += step;
+    }
+
+    return Value(output);
+}
+
+REGISTER_EXPRESSION(range, ExpressionRange::parse);
+const char* ExpressionRange::getOpName() const {
+    return "$range";
+}
+
 /* ------------------------ ExpressionReverseArray ------------------------ */
 
 Value ExpressionReverseArray::evaluateInternal(Variables* vars) const {
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 07ba3f23ff2..0767d70ca64 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1125,6 +1125,12 @@ class ExpressionPow final : public ExpressionFixedArity<ExpressionPow, 2> {
 };
 
 
+class ExpressionRange final : public ExpressionRangedArity<ExpressionRange, 2, 3> {
+    Value evaluateInternal(Variables* vars) const final;
+    const char* getOpName() const final;
+};
+
+
 class ExpressionSecond final : public ExpressionFixedArity<ExpressionSecond, 1> {
 public:
     Value evaluateInternal(Variables* vars) const final;
diff --git a/src/mongo/db/pipeline/expression_test.cpp b/src/mongo/db/pipeline/expression_test.cpp
index a9ccb39c58b..29c3e463690 100644
--- a/src/mongo/db/pipeline/expression_test.cpp
+++ b/src/mongo/db/pipeline/expression_test.cpp
@@ -757,6 +757,41 @@ TEST_F(ExpressionFloorTest, NullArg) {
     assertEvaluates(Value(BSONNULL), Value(BSONNULL));
 }
 
+/* ------------------------ ExpressionRange --------------------------- */
+
+TEST(ExpressionRangeTest, ComputesStandardRange) {
+    assertExpectedResults("$range", {{{Value(0), Value(3)}, Value(BSON_ARRAY(0 << 1 << 2))}});
+}
+
+TEST(ExpressionRangeTest, ComputesRangeWithStep) {
+    assertExpectedResults("$range",
+                          {{{Value(0), Value(6), Value(2)}, Value(BSON_ARRAY(0 << 2 << 4))}});
+}
+
+TEST(ExpressionRangeTest, ComputesReverseRange) {
+    assertExpectedResults("$range",
+                          {{{Value(0), Value(-3), Value(-1)}, Value(BSON_ARRAY(0 << -1 << -2))}});
+}
+
+TEST(ExpressionRangeTest, ComputesRangeWithPositiveAndNegative) {
+    assertExpectedResults("$range",
+                          {{{Value(-2), Value(3)}, Value(BSON_ARRAY(-2 << -1 << 0 << 1 << 2))}});
+}
+
+TEST(ExpressionRangeTest, ComputesEmptyRange) {
+    assertExpectedResults("$range",
+                          {{{Value(-2), Value(3), Value(-1)}, Value(std::vector<Value>())}});
+}
+
+TEST(ExpressionRangeTest, ComputesRangeWithSameStartAndEnd) {
+    assertExpectedResults("$range", {{{Value(20), Value(20)}, Value(std::vector<Value>())}});
+}
+
+TEST(ExpressionRangeTest, ComputesRangeWithLargeNegativeStep) {
+    assertExpectedResults("$range",
+                          {{{Value(3), Value(-5), Value(-3)}, Value(BSON_ARRAY(3 << 0 << -3))}});
+}
+
 /* ------------------------ ExpressionReverseArray -------------------- */
 
 TEST(ExpressionReverseArrayTest, ReversesNormalArray) {