SERVER-10689 Aggregation now supports the switch expression.

8 files changed, 375 insertions, 1 deletions

diff --git a/buildscripts/resmokeconfig/suites/aggregation.yml b/buildscripts/resmokeconfig/suites/aggregation.yml
index c85f2f8d657..aeeaf00531a 100644
--- a/buildscripts/resmokeconfig/suites/aggregation.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation.yml
@@ -3,6 +3,7 @@ selector:
     roots:
     - jstests/aggregation/*.js
     - jstests/aggregation/bugs/*.js
+    - jstests/aggregation/expressions/*.js
 
 executor:
   js_test:
diff --git a/buildscripts/resmokeconfig/suites/aggregation_auth.yml b/buildscripts/resmokeconfig/suites/aggregation_auth.yml
index 65c083c449a..73dcb239d1a 100644
--- a/buildscripts/resmokeconfig/suites/aggregation_auth.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation_auth.yml
@@ -8,6 +8,7 @@ selector:
     roots:
     - jstests/aggregation/*.js
     - jstests/aggregation/bugs/*.js
+    - jstests/aggregation/expressions/*.js
     exclude_files:
     # Skip any tests that run with auth explicitly.
     - jstests/aggregation/*[aA]uth*.js
diff --git a/buildscripts/resmokeconfig/suites/aggregation_ese.yml b/buildscripts/resmokeconfig/suites/aggregation_ese.yml
index b02274361fa..4db6693239a 100644
--- a/buildscripts/resmokeconfig/suites/aggregation_ese.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation_ese.yml
@@ -7,6 +7,7 @@ selector:
     roots:
     - jstests/aggregation/*.js
     - jstests/aggregation/bugs/*.js
+    - jstests/aggregation/expressions/*.js
     - src/mongo/db/modules/*/jstests/aggregation/*.js
     exclude_files:
     # Skip any tests that run with auth explicitly.
diff --git a/buildscripts/resmokeconfig/suites/aggregation_read_concern_majority_passthrough.yml b/buildscripts/resmokeconfig/suites/aggregation_read_concern_majority_passthrough.yml
index 79bb588fdba..842863c23e8 100644
--- a/buildscripts/resmokeconfig/suites/aggregation_read_concern_majority_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation_read_concern_majority_passthrough.yml
@@ -3,6 +3,7 @@ selector:
     roots:
     - jstests/aggregation/*.js
     - jstests/aggregation/bugs/*.js
+    - jstests/aggregation/expressions/*.js
     exclude_files:
     - jstests/aggregation/bugs/server18198.js # Uses a mocked mongo client to test read preference.
     - jstests/aggregation/mongos_slaveok.js # Majority read on secondary requires afterOpTime.
@@ -28,4 +29,4 @@ executor:
           enableTestCommands: 1
       num_nodes: 2
       # Needs to be set for any ephemeral or no-journaling storage engine
-      write_concern_majority_journal_default: false
\ No newline at end of file
+      write_concern_majority_journal_default: false
diff --git a/jstests/aggregation/expressions/switch.js b/jstests/aggregation/expressions/switch.js
new file mode 100644
index 00000000000..9a6dbbb529d
--- /dev/null
+++ b/jstests/aggregation/expressions/switch.js
@@ -0,0 +1,150 @@
+// In SERVER-10689, the $switch expression was introduced. In this file, we test the functionality
+// of the expression.
+
+(function() {
+    "use strict";
+
+    var coll = db.switch;
+    coll.drop();
+
+    // Insert an empty document so that something can flow through the pipeline.
+    coll.insert({});
+
+    // Ensure that a branch is correctly evaluated.
+    var pipeline = {
+        "$project": {
+            "_id": 0,
+            "output": {
+                "$switch": {
+                    "branches": [{"case": {"$eq": [1, 1]}, "then": "one is equal to one!"}],
+                }
+            }
+        }
+    };
+    var res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {"output": "one is equal to one!"});
+
+    // Ensure that the first branch which matches is chosen.
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output": {
+                "$switch": {
+                    "branches": [
+                        {"case": {"$eq": [1, 1]}, "then": "one is equal to one!"},
+                        {"case": {"$eq": [2, 2]}, "then": "two is equal to two!"}
+                    ],
+                }
+            }
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {"output": "one is equal to one!"});
+
+    // Ensure that the default is chosen if no case matches.
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output": {
+                "$switch": {
+                    "branches": [{"case": {"$eq": [1, 2]}, "then": "one is equal to two!"}],
+                    "default": "no case matched."
+                }
+            }
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {"output": "no case matched."});
+
+    // Ensure that nullish values are treated as false when they are a "case", and are null
+    // otherwise.
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output": {
+                "$switch": {
+                    "branches": [{"case": null, "then": "Null was true!"}],
+                    "default": "No case matched."
+                }
+            }
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {"output": "No case matched."});
+
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output": {
+                "$switch": {
+                    "branches": [{"case": "$missingField", "then": "Null was true!"}],
+                    "default": "No case matched."
+                }
+            }
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {"output": "No case matched."});
+
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output":
+                {"$switch": {"branches": [{"case": true, "then": null}], "default": false}}
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {"output": null});
+
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output": {
+                "$switch":
+                    {"branches": [{"case": true, "then": "$missingField"}], "default": false}
+            }
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {});
+
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output":
+                {"$switch": {"branches": [{"case": null, "then": false}], "default": null}}
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {"output": null});
+
+    pipeline = {
+        "$project": {
+            "_id": 0,
+            "output": {
+                "$switch":
+                    {"branches": [{"case": null, "then": false}], "default": "$missingField"}
+            }
+        }
+    };
+    res = coll.aggregate(pipeline).toArray();
+
+    assert.eq(res.length, 1);
+    assert.eq(res[0], {});
+}());
diff --git a/jstests/aggregation/expressions/switch_errors.js b/jstests/aggregation/expressions/switch_errors.js
new file mode 100644
index 00000000000..cf6dc0f4f93
--- /dev/null
+++ b/jstests/aggregation/expressions/switch_errors.js
@@ -0,0 +1,71 @@
+// SERVER-10689 introduced the $switch expression. In this file, we test the error cases of the
+// expression.
+load("jstests/aggregation/extras/utils.js");  // For assertErrorCode.
+
+(function() {
+    "use strict";
+
+    var coll = db.switch;
+    coll.drop();
+
+    var pipeline = {
+        "$project": {"output": {"$switch": "not an object"}}
+    };
+    assertErrorCode(coll, pipeline, 40060, "$switch requires an object as an argument.");
+
+    pipeline = {
+        "$project": {"output": {"$switch": {"branches": "not an array"}}}
+    };
+    assertErrorCode(coll, pipeline, 40061, "$switch requires 'branches' to be an array.");
+
+    pipeline = {
+        "$project": {"output": {"$switch": {"branches": ["not an object"]}}}
+    };
+    assertErrorCode(coll, pipeline, 40062, "$switch requires each branch to be an object.");
+
+    pipeline = {
+        "$project": {"output": {"$switch": {"branches": [{}]}}}
+    };
+    assertErrorCode(coll, pipeline, 40064, "$switch requires each branch have a 'case'.");
+
+    pipeline = {
+        "$project": {
+            "output": {
+                "$switch": {
+                    "branches": [{
+                        "case": 1,
+                    }]
+                }
+            }
+        }
+    };
+    assertErrorCode(coll, pipeline, 40065, "$switch requires each branch have a 'then'.");
+
+    pipeline = {
+        "$project":
+            {"output": {"$switch": {"branches": [{"case": true, "then": false, "badKey": 1}]}}}
+    };
+    assertErrorCode(coll, pipeline, 40063, "$switch found a branch with an unknown argument");
+
+    pipeline = {
+        "$project": {"output": {"$switch": {"notAnArgument": 1}}}
+    };
+    assertErrorCode(coll, pipeline, 40067, "$switch found an unknown argument");
+
+    pipeline = {
+        "$project": {"output": {"$switch": {"branches": []}}}
+    };
+    assertErrorCode(coll, pipeline, 40068, "$switch requires at least one branch");
+
+    pipeline = {
+        "$project": {"output": {"$switch": {}}}
+    };
+    assertErrorCode(coll, pipeline, 40068, "$switch requires at least one branch");
+
+    coll.insert({x: 1});
+    pipeline = {
+        "$project":
+            {"output": {"$switch": {"branches": [{"case": {"$eq": ["$x", 0]}, "then": 1}]}}}
+    };
+    assertErrorCode(coll, pipeline, 40066, "$switch has no default and an input matched no case");
+}());
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 6b4fd5ac6ee..4f1f050478b 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -3264,6 +3264,136 @@ const char* ExpressionSubtract::getOpName() const {
     return "$subtract";
 }
 
+/* ------------------------- ExpressionSwitch ------------------------------ */
+
+REGISTER_EXPRESSION(switch, ExpressionSwitch::parse);
+const char* ExpressionSwitch::getOpName() const {
+    return "$switch";
+}
+
+Value ExpressionSwitch::evaluateInternal(Variables* vars) const {
+    for (auto&& branch : _branches) {
+        Value caseExpression(branch.first->evaluateInternal(vars));
+
+        if (caseExpression.coerceToBool()) {
+            return branch.second->evaluateInternal(vars);
+        }
+    }
+
+    uassert(40066,
+            "$switch could not find a matching branch for an input, and no default was specified.",
+            _default);
+
+    return _default->evaluateInternal(vars);
+}
+
+boost::intrusive_ptr<Expression> ExpressionSwitch::parse(BSONElement expr,
+                                                         const VariablesParseState& vps) {
+    uassert(40060,
+            str::stream() << "$switch requires an object as an argument, found: "
+                          << typeName(expr.type()),
+            expr.type() == Object);
+
+    intrusive_ptr<ExpressionSwitch> expression(new ExpressionSwitch());
+
+    for (auto&& elem : expr.Obj()) {
+        auto field = elem.fieldNameStringData();
+
+        if (field == "branches") {
+            // Parse each branch separately.
+            uassert(40061,
+                    str::stream() << "$switch expected an array for 'branches', found: "
+                                  << typeName(elem.type()),
+                    elem.type() == Array);
+
+            for (auto&& branch : elem.Array()) {
+                uassert(40062,
+                        str::stream() << "$switch expected each branch to be an object, found: "
+                                      << typeName(branch.type()),
+                        branch.type() == Object);
+
+                ExpressionPair branchExpression;
+
+                for (auto&& branchElement : branch.Obj()) {
+                    auto branchField = branchElement.fieldNameStringData();
+
+                    if (branchField == "case") {
+                        branchExpression.first = parseOperand(branchElement, vps);
+                    } else if (branchField == "then") {
+                        branchExpression.second = parseOperand(branchElement, vps);
+                    } else {
+                        uasserted(40063,
+                                  str::stream() << "$switch found an unknown argument to a branch: "
+                                                << branchField);
+                    }
+                }
+
+                uassert(40064,
+                        "$switch requires each branch have a 'case' expression",
+                        branchExpression.first);
+                uassert(40065,
+                        "$switch requires each branch have a 'then' expression.",
+                        branchExpression.second);
+
+                expression->_branches.push_back(branchExpression);
+            }
+        } else if (field == "default") {
+            // Optional, arbitrary expression.
+            expression->_default = parseOperand(elem, vps);
+        } else {
+            uasserted(40067, str::stream() << "$switch found an unknown argument: " << field);
+        }
+    }
+
+    uassert(40068, "$switch requires at least one branch.", !expression->_branches.empty());
+
+    return expression;
+}
+
+void ExpressionSwitch::addDependencies(DepsTracker* deps, std::vector<std::string>* path) const {
+    for (auto&& branch : _branches) {
+        branch.first->addDependencies(deps, path);
+        branch.second->addDependencies(deps, path);
+    }
+
+    if (_default) {
+        _default->addDependencies(deps, path);
+    }
+}
+
+boost::intrusive_ptr<Expression> ExpressionSwitch::optimize() {
+    if (_default) {
+        _default = _default->optimize();
+    }
+
+    std::transform(_branches.begin(),
+                   _branches.end(),
+                   _branches.begin(),
+                   [](ExpressionPair branch) -> ExpressionPair {
+                       return {branch.first->optimize(), branch.second->optimize()};
+                   });
+
+    return this;
+}
+
+Value ExpressionSwitch::serialize(bool explain) const {
+    std::vector<Value> serializedBranches;
+    serializedBranches.reserve(_branches.size());
+
+    for (auto&& branch : _branches) {
+        serializedBranches.push_back(Value(Document{{"case", branch.first->serialize(explain)},
+                                                    {"then", branch.second->serialize(explain)}}));
+    }
+
+    if (_default) {
+        return Value(Document{{"$switch",
+                               Document{{"branches", Value(serializedBranches)},
+                                        {"default", _default->serialize(explain)}}}});
+    }
+
+    return Value(Document{{"$switch", Document{{"branches", Value(serializedBranches)}}}});
+}
+
 /* ------------------------- ExpressionToLower ----------------------------- */
 
 Value ExpressionToLower::evaluateInternal(Variables* vars) const {
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 0939bcbcdc7..a5c78babce1 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1278,6 +1278,25 @@ public:
 };
 
 
+class ExpressionSwitch final : public ExpressionFixedArity<ExpressionSwitch, 1> {
+public:
+    void addDependencies(DepsTracker* deps, std::vector<std::string>* path = nullptr) const final;
+    Value evaluateInternal(Variables* vars) const final;
+    boost::intrusive_ptr<Expression> optimize() final;
+    static boost::intrusive_ptr<Expression> parse(BSONElement expr,
+                                                  const VariablesParseState& vpsIn);
+    Value serialize(bool explain) const final;
+    const char* getOpName() const final;
+
+private:
+    using ExpressionPair =
+        std::pair<boost::intrusive_ptr<Expression>, boost::intrusive_ptr<Expression>>;
+
+    boost::intrusive_ptr<Expression> _default;
+    std::vector<ExpressionPair> _branches;
+};
+
+
 class ExpressionToLower final : public ExpressionFixedArity<ExpressionToLower, 1> {
 public:
     Value evaluateInternal(Variables* vars) const final;