SERVER-30967 Add $binarySize expression

5 files changed, 129 insertions, 7 deletions

diff --git a/jstests/aggregation/expressions/binarySize.js b/jstests/aggregation/expressions/binarySize.js
new file mode 100644
index 00000000000..076498b8f63
--- /dev/null
+++ b/jstests/aggregation/expressions/binarySize.js
@@ -0,0 +1,44 @@
+/**
+ * Test the $binarySize expression.
+ */
+(function() {
+"use strict";
+load("jstests/aggregation/extras/utils.js");
+
+const coll = db.expression_binarySize;
+coll.drop();
+
+assert.commandWorked(coll.insert([
+    {_id: 0, x: ""},
+    {_id: 1, x: "abc"},
+    {_id: 2, x: "ab\0c"},
+    {_id: 3, x: "abc\0"},
+    {_id: 4, x: BinData(0, "")},
+    {_id: 5, x: BinData(0, "1234")},
+    {_id: 6, x: null},
+    {_id: 7},
+]));
+
+const result =
+    coll.aggregate([{$sort: {_id: 1}}, {$addFields: {s: {$binarySize: "$x"}}}]).toArray();
+assert.eq(result, [
+
+    {_id: 0, x: "", s: 0},
+    {_id: 1, x: "abc", s: 3},
+    // Javascript strings and BSON strings can contain '\0', so both of these have length 4.
+    {_id: 2, x: "ab\0c", s: 4},
+    {_id: 3, x: "abc\0", s: 4},
+
+    {_id: 4, x: BinData(0, ""), s: 0},
+    // The mongo shell BinData constructor takes base64, so "1234" encodes 3 bytes.
+    {_id: 5, x: BinData(0, "1234"), s: 3},
+
+    // $binarySize also accepts nullish values, and returns null.
+    {_id: 6, x: null, s: null},
+    {_id: 7, s: null},
+]);
+
+// $binarySize only accepts strings and BinData.
+assert.commandWorked(coll.insert({x: 42}));
+assertErrorCode(coll, {$project: {s: {$binarySize: "$x"}}}, 51276);
+}());
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 19658f2e005..ee960534f0e 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -4416,20 +4416,26 @@ const char* ExpressionSubstrCP::getOpName() const {
 
 /* ----------------------- ExpressionStrLenBytes ------------------------- */
 
+namespace {
+Value strLenBytes(StringData str) {
+    size_t strLen = str.size();
+
+    uassert(34470,
+            "string length could not be represented as an int.",
+            strLen <= std::numeric_limits<int>::max());
+    return Value(static_cast<int>(strLen));
+}
+}  // namespace
+
 Value ExpressionStrLenBytes::evaluate(const Document& root, Variables* variables) const {
     Value str(_children[0]->evaluate(root, variables));
 
     uassert(34473,
             str::stream() << "$strLenBytes requires a string argument, found: "
                           << typeName(str.getType()),
-            str.getType() == String);
-
-    size_t strLen = str.getString().size();
+            str.getType() == BSONType::String);
 
-    uassert(34470,
-            "string length could not be represented as an int.",
-            strLen <= std::numeric_limits<int>::max());
-    return Value(static_cast<int>(strLen));
+    return strLenBytes(str.getStringData());
 }
 
 REGISTER_EXPRESSION(strLenBytes, ExpressionStrLenBytes::parse);
@@ -4437,6 +4443,32 @@ const char* ExpressionStrLenBytes::getOpName() const {
     return "$strLenBytes";
 }
 
+/* -------------------------- ExpressionBinarySize ------------------------------ */
+
+Value ExpressionBinarySize::evaluate(const Document& root, Variables* variables) const {
+    Value arg = _children[0]->evaluate(root, variables);
+    if (arg.nullish()) {
+        return Value(BSONNULL);
+    }
+
+    uassert(51276,
+            str::stream() << "$binarySize requires a string or BinData argument, found: "
+                          << typeName(arg.getType()),
+            arg.getType() == BSONType::BinData || arg.getType() == BSONType::String);
+
+    if (arg.getType() == BSONType::String) {
+        return strLenBytes(arg.getStringData());
+    }
+
+    BSONBinData binData = arg.getBinData();
+    return Value(binData.length);
+}
+
+REGISTER_EXPRESSION(binarySize, ExpressionBinarySize::parse);
+const char* ExpressionBinarySize::getOpName() const {
+    return "$binarySize";
+}
+
 /* ----------------------- ExpressionStrLenCP ------------------------- */
 
 Value ExpressionStrLenCP::evaluate(const Document& root, Variables* variables) const {
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index 21fff1d3ce9..c2dae92cffb 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -2202,6 +2202,20 @@ public:
 };
 
 
+class ExpressionBinarySize final : public ExpressionFixedArity<ExpressionBinarySize, 1> {
+public:
+    ExpressionBinarySize(const boost::intrusive_ptr<ExpressionContext>& expCtx)
+        : ExpressionFixedArity<ExpressionBinarySize, 1>(expCtx) {}
+
+    Value evaluate(const Document& root, Variables* variables) const final;
+    const char* getOpName() const final;
+
+    void acceptVisitor(ExpressionVisitor* visitor) final {
+        return visitor->visit(this);
+    }
+};
+
+
 class ExpressionStrLenCP final : public ExpressionFixedArity<ExpressionStrLenCP, 1> {
 public:
     explicit ExpressionStrLenCP(const boost::intrusive_ptr<ExpressionContext>& expCtx)
diff --git a/src/mongo/db/pipeline/expression_nary_test.cpp b/src/mongo/db/pipeline/expression_nary_test.cpp
index 1506ee61c68..b96220378ce 100644
--- a/src/mongo/db/pipeline/expression_nary_test.cpp
+++ b/src/mongo/db/pipeline/expression_nary_test.cpp
@@ -1188,5 +1188,35 @@ TEST_F(ExpressionRoundTwoArgTest, NullArg2) {
     assertEval(BSONNULL, 1, BSONNULL);
 }
 
+/* ------------------------- ExpressionBinarySize -------------------------- */
+
+class ExpressionBinarySizeTest : public ExpressionNaryTestOneArg {
+public:
+    void assertEval(ImplicitValue input, ImplicitValue output) {
+        intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest());
+        _expr = new ExpressionBinarySize(expCtx);
+        ExpressionNaryTestOneArg::assertEvaluates(input, output);
+    }
+};
+
+TEST_F(ExpressionBinarySizeTest, HandlesStrings) {
+    assertEval("abc"_sd, 3);
+    assertEval(""_sd, 0);
+    assertEval("ab\0c"_sd, 4);
+    assertEval("abc\0"_sd, 4);
+}
+
+TEST_F(ExpressionBinarySizeTest, HandlesBinData) {
+    assertEval(BSONBinData("abc", 3, BinDataGeneral), 3);
+    assertEval(BSONBinData("", 0, BinDataGeneral), 0);
+    assertEval(BSONBinData("ab\0c", 4, BinDataGeneral), 4);
+    assertEval(BSONBinData("abc\0", 4, BinDataGeneral), 4);
+}
+
+TEST_F(ExpressionBinarySizeTest, HandlesNullish) {
+    assertEval(BSONNULL, BSONNULL);
+    assertEval(BSONUndefined, BSONNULL);
+}
+
 }  // anonymous namespace
 }  // namespace ExpressionTests
diff --git a/src/mongo/db/pipeline/expression_visitor.h b/src/mongo/db/pipeline/expression_visitor.h
index f4ccd264aa8..8269da1e1f4 100644
--- a/src/mongo/db/pipeline/expression_visitor.h
+++ b/src/mongo/db/pipeline/expression_visitor.h
@@ -106,6 +106,7 @@ class ExpressionStrcasecmp;
 class ExpressionSubstrBytes;
 class ExpressionSubstrCP;
 class ExpressionStrLenBytes;
+class ExpressionBinarySize;
 class ExpressionStrLenCP;
 class ExpressionSubtract;
 class ExpressionSwitch;
@@ -231,6 +232,7 @@ public:
     virtual void visit(ExpressionSubstrBytes*) = 0;
     virtual void visit(ExpressionSubstrCP*) = 0;
     virtual void visit(ExpressionStrLenBytes*) = 0;
+    virtual void visit(ExpressionBinarySize*) = 0;
     virtual void visit(ExpressionStrLenCP*) = 0;
     virtual void visit(ExpressionSubtract*) = 0;
     virtual void visit(ExpressionSwitch*) = 0;