SERVER-27614 add $$REMOVE agg system variable

8 files changed, 290 insertions, 28 deletions

diff --git a/buildscripts/resmokeconfig/suites/aggregation.yml b/buildscripts/resmokeconfig/suites/aggregation.yml
index c7a8c6e64be..dcfea8ecf3f 100644
--- a/buildscripts/resmokeconfig/suites/aggregation.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation.yml
@@ -5,6 +5,7 @@ selector:
     - jstests/aggregation/bugs/*.js
     - jstests/aggregation/expressions/*.js
     - jstests/aggregation/sources/*/*.js
+    - jstests/aggregation/variables/*.js
 
 executor:
   js_test:
diff --git a/jstests/aggregation/variables/remove_system_variable.js b/jstests/aggregation/variables/remove_system_variable.js
new file mode 100644
index 00000000000..803c826af4f
--- /dev/null
+++ b/jstests/aggregation/variables/remove_system_variable.js
@@ -0,0 +1,87 @@
+/**
+ * Tests for the $$REMOVE system variable.
+ */
+(function() {
+    "use strict";
+
+    let coll = db[jsTest.name()];
+    coll.drop();
+
+    assert.writeOK(coll.insert({_id: 1, a: 2, b: 3}));
+    assert.writeOK(coll.insert({_id: 2, a: 3, b: 4}));
+    assert.writeOK(coll.insert({_id: 3, a: {b: 98, c: 99}}));
+
+    let projectStage = {
+        $project: {_id: 0, a: 1, b: {$cond: {if: {$eq: ["$b", 4]}, then: "$$REMOVE", else: "$b"}}}
+    };
+
+    // Test that we can conditionally remove a field in $project.
+    assert.eq([{a: 2, b: 3}], coll.aggregate([{$match: {_id: 1}}, projectStage]).toArray());
+    assert.eq([{a: 3}], coll.aggregate([{$match: {_id: 2}}, projectStage]).toArray());
+
+    // Test removal of a nested field, using $project.
+    assert.eq([{a: {b: 98}}],
+              coll.aggregate([{$match: {_id: 3}}, {$project: {_id: 0, "a.b": 1}}]).toArray());
+    assert.eq(
+        [{a: {}}],
+        coll.aggregate([{$match: {_id: 3}}, {$project: {_id: 0, "a.b": "$$REMOVE"}}]).toArray());
+    assert.eq(
+        [{a: {}}],
+        coll.aggregate([{$match: {_id: 3}}, {$project: {_id: 0, a: {b: "$$REMOVE"}}}]).toArray());
+
+    // Test removal of a nested field, using $addFields.
+    assert.eq([{_id: 3, a: {c: 99}}],
+              coll.aggregate([{$match: {_id: 3}}, {$addFields: {"a.b": "$$REMOVE"}}]).toArray());
+
+    // Test that any field path following "$$REMOVE" also evaluates to missing.
+    assert.eq([{_id: 3}],
+              coll.aggregate([{$match: {_id: 3}}, {$addFields: {"a": "$$REMOVE.a.c"}}]).toArray());
+
+    // Test that $$REMOVE can be used together with user-defined variables in a $let.
+    assert.eq([{a: {b: 3, d: 4}}],
+              coll.aggregate([
+                      {$match: {_id: 3}},
+                      {
+                        $project: {
+                            _id: 0,
+                            a: {
+                                $let: {
+                                    vars: {bar: 3, foo: 4},
+                                    in : {b: "$$bar", c: "$$REMOVE", d: "$$foo"}
+                                }
+                            }
+                        }
+                      }
+                  ])
+                  .toArray());
+
+    // Test that $$REMOVE cannot be assigned in a $let.
+    assert.commandFailedWithCode(db.runCommand({
+        aggregate: coll.getName(),
+        cursor: {},
+        pipeline: [
+            {$match: {_id: 3}},
+            {$project: {_id: 0, a: {$let: {vars: {"REMOVE": 3}, in : {b: "$$REMOVE", c: 2}}}}}
+        ]
+    }),
+                                 16867);
+
+    // Test that $$REMOVE, $$CURRENT, $$ROOT, and user-defined variables can all be used together.
+    assert.eq(
+        [{a: {b: 3, d: {_id: 1, a: 2, b: 3}, e: {_id: 1, a: 2, b: 3}}}],
+        coll.aggregate([
+                {$match: {_id: 1}},
+                {
+                  $project: {
+                      _id: 0,
+                      a: {
+                          $let: {
+                              vars: {myVar: 3},
+                              in : {b: "$$myVar", c: "$$REMOVE", d: "$$ROOT", e: "$$CURRENT"}
+                          }
+                      }
+                  }
+                }
+            ])
+            .toArray());
+}());
diff --git a/src/mongo/db/pipeline/document_source_add_fields_test.cpp b/src/mongo/db/pipeline/document_source_add_fields_test.cpp
index 53ff48ae906..be29dea762b 100644
--- a/src/mongo/db/pipeline/document_source_add_fields_test.cpp
+++ b/src/mongo/db/pipeline/document_source_add_fields_test.cpp
@@ -152,5 +152,47 @@ TEST_F(AddFieldsTest, ShouldPropagatePauses) {
     ASSERT_TRUE(addFields->getNext().isEOF());
 }
 
+TEST_F(AddFieldsTest, AddFieldsWithRemoveSystemVariableDoesNotAddField) {
+    auto addFields = DocumentSourceAddFields::create(BSON("fieldToAdd"
+                                                          << "$$REMOVE"),
+                                                     getExpCtx());
+    auto mock = DocumentSourceMock::create(Document{{"existingField", 1}});
+    addFields->setSource(mock.get());
+
+    auto next = addFields->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    Document expected{{"existingField", 1}};
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+    ASSERT_TRUE(addFields->getNext().isEOF());
+}
+
+TEST_F(AddFieldsTest, AddFieldsWithRootSystemVariableAddsRootAsSubDoc) {
+    auto addFields = DocumentSourceAddFields::create(BSON("b"
+                                                          << "$$ROOT"),
+                                                     getExpCtx());
+    auto mock = DocumentSourceMock::create(Document{{"a", 1}});
+    addFields->setSource(mock.get());
+
+    auto next = addFields->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    Document expected{{"a", 1}, {"b", Document{{"a", 1}}}};
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+    ASSERT_TRUE(addFields->getNext().isEOF());
+}
+
+TEST_F(AddFieldsTest, AddFieldsWithCurrentSystemVariableAddsRootAsSubDoc) {
+    auto addFields = DocumentSourceAddFields::create(BSON("b"
+                                                          << "$$CURRENT"),
+                                                     getExpCtx());
+    auto mock = DocumentSourceMock::create(Document{{"a", 1}});
+    addFields->setSource(mock.get());
+
+    auto next = addFields->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    Document expected{{"a", 1}, {"b", Document{{"a", 1}}}};
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+    ASSERT_TRUE(addFields->getNext().isEOF());
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_project_test.cpp b/src/mongo/db/pipeline/document_source_project_test.cpp
index 1247de01579..60fe7a5988f 100644
--- a/src/mongo/db/pipeline/document_source_project_test.cpp
+++ b/src/mongo/db/pipeline/document_source_project_test.cpp
@@ -38,6 +38,7 @@
 #include "mongo/db/pipeline/dependencies.h"
 #include "mongo/db/pipeline/document_source_mock.h"
 #include "mongo/db/pipeline/document_source_project.h"
+#include "mongo/db/pipeline/document_value_test_util.h"
 #include "mongo/db/pipeline/value.h"
 #include "mongo/unittest/unittest.h"
 
@@ -242,5 +243,23 @@ TEST_F(ProjectStageTest, ExclusionProjectionReportsExcludedPathsWithIdExclusion)
     ASSERT_EQUALS(1U, modifiedPaths.paths.count("e.f.g"));
 }
 
+TEST_F(ProjectStageTest, CanUseRemoveSystemVariableToConditionallyExcludeProjectedField) {
+    auto project = DocumentSourceProject::create(
+        fromjson("{a: 1, b: {$cond: [{$eq: ['$b', 4]}, '$$REMOVE', '$b']}}"), getExpCtx());
+    auto source = DocumentSourceMock::create({"{a: 2, b: 2}", "{a: 3, b: 4}"});
+    project->setSource(source.get());
+    auto next = project->getNext();
+    ASSERT(next.isAdvanced());
+    Document expected{{"a", 2}, {"b", 2}};
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+
+    next = project->getNext();
+    ASSERT(next.isAdvanced());
+    expected = Document{{"a", 3}};
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+
+    ASSERT(project->getNext().isEOF());
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_replace_root_test.cpp b/src/mongo/db/pipeline/document_source_replace_root_test.cpp
index ae8da121996..51411531022 100644
--- a/src/mongo/db/pipeline/document_source_replace_root_test.cpp
+++ b/src/mongo/db/pipeline/document_source_replace_root_test.cpp
@@ -279,6 +279,16 @@ TEST_F(ReplaceRootBasics, ReplaceRootModifiesAllFields) {
     ASSERT_EQUALS(0U, modifiedPaths.paths.size());
 }
 
+TEST_F(ReplaceRootBasics, ReplaceRootWithRemoveSystemVariableThrows) {
+    auto replaceRoot = createReplaceRoot(BSON("newRoot"
+                                              << "$$REMOVE"));
+    Document inputDoc = Document{{"b", 2}};
+    auto mock = DocumentSourceMock::create({inputDoc});
+    replaceRoot->setSource(mock.get());
+
+    ASSERT_THROWS_CODE(replaceRoot->getNext(), UserException, 40228);
+}
+
 /**
  * Fixture to test error cases of initializing the $replaceRoot stage.
  */
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index d92e4c68bd9..30cafd8ba23 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -61,9 +61,16 @@ using std::vector;
 
 /// Helper function to easily wrap constants with $const.
 static Value serializeConstant(Value val) {
+    if (val.missing()) {
+        return Value("$$REMOVE"_sd);
+    }
+
     return Value(DOC("$const" << val));
 }
 
+const StringMap<Variables::Id> Variables::kBuiltinVarNameToId = {{"ROOT"_sd, kRootId},
+                                                                 {"REMOVE"_sd, kRemoveId}};
+
 void Variables::uassertValidNameForUserWrite(StringData varName) {
     // System variables users allowed to write to (currently just one)
     if (varName == "CURRENT") {
@@ -124,26 +131,37 @@ void Variables::uassertValidNameForUserRead(StringData varName) {
 }
 
 void Variables::setValue(Id id, const Value& value) {
-    massert(17199, "can't use Variables::setValue to set ROOT", id != ROOT_ID);
-
-    verify(id < _numVars);
-    _rest[id] = value;
+    uassert(17199, "can't use Variables::setValue to set a reserved builtin variable", id >= 0);
+    auto varIndex = static_cast<size_t>(id);
+    invariant(varIndex < _numVars);
+    _rest[varIndex] = value;
 }
 
 Value Variables::getValue(Id id) const {
-    if (id == ROOT_ID)
-        return Value(_root);
+    if (id < 0) {
+        // This is a reserved id for a builtin variable.
+        switch (id) {
+            case Variables::kRootId:
+                return Value(_root);
+            case Variables::kRemoveId:
+                return Value();
+            default:
+                MONGO_UNREACHABLE;
+        }
+    }
 
-    verify(id < _numVars);
-    return _rest[id];
+    auto varIndex = static_cast<size_t>(id);
+    invariant(varIndex < _numVars);
+    return _rest[varIndex];
 }
 
 Document Variables::getDocument(Id id) const {
-    if (id == ROOT_ID)
+    if (id == Variables::kRootId) {
+        // For the common case of ROOT, avoid round-tripping through Value.
         return _root;
+    }
 
-    verify(id < _numVars);
-    const Value var = _rest[id];
+    const Value var = getValue(id);
     if (var.getType() == Object)
         return var.getDocument();
 
@@ -151,24 +169,34 @@ Document Variables::getDocument(Id id) const {
 }
 
 Variables::Id VariablesParseState::defineVariable(StringData name) {
-    // caller should have validated before hand by using Variables::uassertValidNameForUserWrite
-    massert(17275, "Can't redefine ROOT", name != "ROOT");
+    // Caller should have validated before hand by using Variables::uassertValidNameForUserWrite.
+    massert(17275,
+            "Can't redefine a non-user-writable variable",
+            Variables::kBuiltinVarNameToId.find(name) == Variables::kBuiltinVarNameToId.end());
 
     Variables::Id id = _idGenerator->generateId();
+    invariant(id >= 0);
     _variables[name] = id;
     return id;
 }
 
 Variables::Id VariablesParseState::getVariable(StringData name) const {
-    StringMap<Variables::Id>::const_iterator it = _variables.find(name);
-    if (it != _variables.end())
+    auto it = _variables.find(name);
+    if (it != _variables.end()) {
+        // Found a user-defined variable.
         return it->second;
+    }
 
-    uassert(17276,
-            str::stream() << "Use of undefined variable: " << name,
-            name == "ROOT" || name == "CURRENT");
+    it = Variables::kBuiltinVarNameToId.find(name);
+    if (it != Variables::kBuiltinVarNameToId.end()) {
+        // This is a builtin variable.
+        return it->second;
+    }
 
-    return Variables::ROOT_ID;
+    // If we didn't find either a user-defined or builtin variable, then we reject everything other
+    // than CURRENT. If this is CURRENT, then we treat it as equivalent to ROOT.
+    uassert(17276, str::stream() << "Use of undefined variable: " << name, name == "CURRENT");
+    return Variables::kRootId;
 }
 
 /* --------------------------- Expression ------------------------------ */
@@ -1419,7 +1447,7 @@ Value ExpressionObject::serialize(bool explain) const {
 // this is the old deprecated version only used by tests not using variables
 intrusive_ptr<ExpressionFieldPath> ExpressionFieldPath::create(
     const boost::intrusive_ptr<ExpressionContext>& expCtx, const string& fieldPath) {
-    return new ExpressionFieldPath(expCtx, "CURRENT." + fieldPath, Variables::ROOT_ID);
+    return new ExpressionFieldPath(expCtx, "CURRENT." + fieldPath, Variables::kRootId);
 }
 
 // this is the new version that supports every syntax
@@ -1454,12 +1482,16 @@ ExpressionFieldPath::ExpressionFieldPath(const boost::intrusive_ptr<ExpressionCo
     : Expression(expCtx), _fieldPath(theFieldPath), _variable(variable) {}
 
 intrusive_ptr<Expression> ExpressionFieldPath::optimize() {
-    /* nothing can be done for these */
+    if (_variable == Variables::kRemoveId) {
+        // The REMOVE system variable optimizes to a constant missing value.
+        return ExpressionConstant::create(getExpressionContext(), Value());
+    }
+
     return intrusive_ptr<Expression>(this);
 }
 
 void ExpressionFieldPath::addDependencies(DepsTracker* deps) const {
-    if (_variable == Variables::ROOT_ID) {  // includes CURRENT when it is equivalent to ROOT.
+    if (_variable == Variables::kRootId) {  // includes CURRENT when it is equivalent to ROOT.
         if (_fieldPath.getPathLength() == 1) {
             deps->needWholeDocument = true;  // need full doc if just "$$ROOT"
         } else {
@@ -1511,7 +1543,7 @@ Value ExpressionFieldPath::evaluateInternal(Variables* vars) const {
     if (_fieldPath.getPathLength() == 1)  // get the whole variable
         return vars->getValue(_variable);
 
-    if (_variable == Variables::ROOT_ID) {
+    if (_variable == Variables::kRootId) {
         // ROOT is always a document so use optimized code path
         return evaluatePath(1, vars->getRoot());
     }
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 0e8185ef3db..e876c5e7079 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -72,11 +72,15 @@ class Variables {
 
 public:
     /**
-     * Each unique variable is assigned a unique id of this type
+     * Each unique variable is assigned a unique id of this type. Negative ids are reserved for
+     * system variables and non-negative ids are allocated for user variables.
      */
-    typedef size_t Id;
+    typedef int64_t Id;
 
-    // This is only for expressions that use no variables (even ROOT).
+    /**
+     * Constructs a placeholder for expressions that use no variables (even builtins like ROOT or
+     * REMOVE).
+     */
     Variables() : _numVars(0) {}
 
     explicit Variables(size_t numVars, const Document& root = Document())
@@ -85,7 +89,12 @@ public:
     static void uassertValidNameForUserWrite(StringData varName);
     static void uassertValidNameForUserRead(StringData varName);
 
-    static const Id ROOT_ID = Id(-1);
+    // Ids for builtin variables.
+    static constexpr Id kRootId = Id(-1);
+    static constexpr Id kRemoveId = Id(-2);
+
+    // Map from builtin var name to reserved id number.
+    static const StringMap<Id> kBuiltinVarNameToId;
 
     /**
      * Use this instead of setValue for setting ROOT
@@ -100,11 +109,19 @@ public:
         return _root;
     }
 
+    /**
+     * Sets the value of a user-defined variable. Illegal to use with the reserved builtin variables
+     * defined above.
+     */
     void setValue(Id id, const Value& value);
+
+    /**
+     * Gets the value of a user-defined or system variable.
+     */
     Value getValue(Id id) const;
 
     /**
-     * returns Document() for non-document values.
+     * Returns Document() for non-document values, but otherwise identical to getValue().
      */
     Document getDocument(Id id) const;
 
diff --git a/src/mongo/db/pipeline/expression_test.cpp b/src/mongo/db/pipeline/expression_test.cpp
index 315a2faeac7..f08617396bb 100644
--- a/src/mongo/db/pipeline/expression_test.cpp
+++ b/src/mongo/db/pipeline/expression_test.cpp
@@ -2074,6 +2074,18 @@ private:
     }
 };
 
+TEST(ExpressionConstantTest, ConstantOfValueMissingRemovesField) {
+    intrusive_ptr<Expression> expression = ExpressionConstant::create(nullptr, Value());
+    assertBinaryEqual(BSONObj(), toBson(expression->evaluate(Document{{"foo", Value("bar"_sd)}})));
+}
+
+TEST(ExpressionConstantTest, ConstantOfValueMissingSerializesToRemoveSystemVar) {
+    intrusive_ptr<Expression> expression = ExpressionConstant::create(nullptr, Value());
+    assertBinaryEqual(BSON("field"
+                           << "$$REMOVE"),
+                      BSON("field" << expression->serialize(false)));
+}
+
 }  // namespace Constant
 
 TEST(ExpressionFromAccumulators, Avg) {
@@ -3836,6 +3848,48 @@ TEST(ExpressionTypeTest, WithMaxKeyValue) {
 
 }  // namespace Type
 
+namespace BuiltinRemoveVariable {
+
+TEST(BuiltinRemoveVariableTest, TypeOfRemoveIsMissing) {
+    assertExpectedResults("$type", {{{Value("$$REMOVE"_sd)}, Value("missing"_sd)}});
+}
+
+TEST(BuiltinRemoveVariableTest, LiteralEscapesRemoveVar) {
+    assertExpectedResults(
+        "$literal", {{{Value("$$REMOVE"_sd)}, Value(std::vector<Value>{Value("$$REMOVE"_sd)})}});
+}
+
+TEST(BuiltinRemoveVariableTest, RemoveSerializesCorrectly) {
+    VariablesIdGenerator idGenerator{};
+    VariablesParseState vps{&idGenerator};
+    auto expression = ExpressionFieldPath::parse(nullptr, "$$REMOVE", vps);
+    ASSERT_BSONOBJ_EQ(BSON("foo"
+                           << "$$REMOVE"),
+                      BSON("foo" << expression->serialize(false)));
+}
+
+TEST(BuiltinRemoveVariableTest, RemoveSerializesCorrectlyWithTrailingPath) {
+    VariablesIdGenerator idGenerator{};
+    VariablesParseState vps{&idGenerator};
+    auto expression = ExpressionFieldPath::parse(nullptr, "$$REMOVE.a.b", vps);
+    ASSERT_BSONOBJ_EQ(BSON("foo"
+                           << "$$REMOVE.a.b"),
+                      BSON("foo" << expression->serialize(false)));
+}
+
+TEST(BuiltinRemoveVariableTest, RemoveSerializesCorrectlyAfterOptimization) {
+    VariablesIdGenerator idGenerator{};
+    VariablesParseState vps{&idGenerator};
+    auto expression = ExpressionFieldPath::parse(nullptr, "$$REMOVE.a.b", vps);
+    auto optimizedExpression = expression->optimize();
+    ASSERT(dynamic_cast<ExpressionConstant*>(optimizedExpression.get()));
+    ASSERT_BSONOBJ_EQ(BSON("foo"
+                           << "$$REMOVE"),
+                      BSON("foo" << optimizedExpression->serialize(false)));
+}
+
+}  // namespace BuiltinRemoveVariable
+
 namespace ToLower {
 
 class ExpectedResultBase {