SERVER-64442 Implement $size for new optimizer

3 files changed, 43 insertions, 2 deletions

diff --git a/jstests/cqf/array_size.js b/jstests/cqf/array_size.js
new file mode 100644
index 00000000000..6b3f16f9b5b
--- /dev/null
+++ b/jstests/cqf/array_size.js
@@ -0,0 +1,24 @@
+(function() {
+"use strict";
+
+load("jstests/libs/optimizer_utils.js");  // For checkCascadesOptimizerEnabled.
+if (!checkCascadesOptimizerEnabled(db)) {
+    jsTestLog("Skipping test because the optimizer is not enabled");
+    return;
+}
+
+const t = db.cqf_array_size;
+t.drop();
+
+assert.commandWorked(t.insert({a: [1, 2, 3, 4]}));
+assert.commandWorked(t.insert({a: [2, 3, 4]}));
+assert.commandWorked(t.insert({a: [2, 2]}));
+assert.commandWorked(t.insert({a: 2}));
+assert.commandWorked(t.insert({a: [1, 3]}));
+
+let res = t.explain("executionStats").aggregate([{$match: {a: {$size: 2}}}]);
+assert.eq(2, res.executionStats.nReturned);
+
+res = t.explain("executionStats").aggregate([{$project: {b: {$size: '$a'}}}, {$match: {b: 2}}]);
+assert.eq(2, res.executionStats.nReturned);
+}());
diff --git a/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp b/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp
index ff6dc5e3b45..f44dc146080 100644
--- a/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp
+++ b/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp
@@ -394,7 +394,7 @@ public:
     }
 
     void visit(const ExpressionSize* expr) override final {
-        unsupportedExpression(expr->getOpName());
+        pushSingleArgFunctionFromTop("getArraySize");
     }
 
     void visit(const ExpressionReverseArray* expr) override final {
diff --git a/src/mongo/db/pipeline/abt/match_expression_visitor.cpp b/src/mongo/db/pipeline/abt/match_expression_visitor.cpp
index 6c33e2d66e3..fb943b9144a 100644
--- a/src/mongo/db/pipeline/abt/match_expression_visitor.cpp
+++ b/src/mongo/db/pipeline/abt/match_expression_visitor.cpp
@@ -272,7 +272,24 @@ public:
     }
 
     void visit(const SizeMatchExpression* expr) override {
-        unsupportedExpression(expr);
+        const std::string lambdaProjName = _prefixId.getNextId("lambda_sizeMatch");
+        ABT result = make<PathLambda>(make<LambdaAbstraction>(
+            lambdaProjName,
+            make<BinaryOp>(
+                Operations::Eq,
+                make<FunctionCall>("getArraySize", makeSeq(make<Variable>(lambdaProjName))),
+                Constant::int64(expr->getData()))));
+
+        if (!expr->path().empty()) {
+            // No traverse.
+            result = translateFieldPath(
+                FieldPath(expr->path().toString()),
+                std::move(result),
+                [](const std::string& fieldName, const bool /*isLastElement*/, ABT input) {
+                    return make<PathGet>(fieldName, std::move(input));
+                });
+        }
+        _ctx.push(std::move(result));
     }
 
     void visit(const TextMatchExpression* expr) override {