SERVER-50762 [SBE] SizeMatchExpression can produce wrong results in some cases involving nesting

2 files changed, 24 insertions, 11 deletions

diff --git a/jstests/core/find_size.js b/jstests/core/find_size.js
index 14949f0f477..9b0864be4c7 100644
--- a/jstests/core/find_size.js
+++ b/jstests/core/find_size.js
@@ -1,5 +1,4 @@
 // Tests the behavior of $size for match expressions.
-
 (function() {
 "use strict";
 
@@ -84,4 +83,10 @@ badInputs.forEach(function(input) {
     assert.commandFailed(db.runCommand({find: coll.getName(), filter: {"a": {$size: input}}}),
                          "$size argument " + input + " should have failed");
 });
+
+assert(coll.drop());
+
+// Check nested arrays.
+assert.commandWorked(coll.insert([{a: 1, b: "foo"}, {a: 1, b: [7, 8, 9]}]));
+assert.eq(1, coll.find({$or: [{$and: [{b: {$size: 3}}, {a: 1}]}, {a: 0}]}).itcount());
 }());
diff --git a/src/mongo/db/query/sbe_stage_builder_filter.cpp b/src/mongo/db/query/sbe_stage_builder_filter.cpp
index a854d78e1b3..161115b05b5 100644
--- a/src/mongo/db/query/sbe_stage_builder_filter.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_filter.cpp
@@ -334,14 +334,6 @@ void generateArraySize(MatchExpressionVisitorContext* context,
 
     auto makePredicate = [&](sbe::value::SlotId inputSlot,
                              EvalStage inputStage) -> EvalExprStagePair {
-        // Before generating a final leaf traverse stage, check that the thing we are about
-        // traverse is indeed an array.
-        auto fromBranch =
-            makeFilter<false>(std::move(inputStage),
-                              sbe::makeE<sbe::EFunction>(
-                                  "isArray", sbe::makeEs(sbe::makeE<sbe::EVariable>(inputSlot))),
-                              context->planNodeId);
-
         // Generate a traverse that projects the integer value 1 for each element in the array and
         // then sums up the 1's, resulting in the count of elements in the array.
         auto innerSlot = context->slotIdGenerator->generate();
@@ -354,7 +346,7 @@ void generateArraySize(MatchExpressionVisitorContext* context,
 
         auto traverseSlot = context->slotIdGenerator->generate();
         auto traverseStage =
-            makeTraverse(std::move(fromBranch),
+            makeTraverse(EvalStage{},
                          std::move(innerBranch),
                          inputSlot,
                          traverseSlot,
@@ -380,7 +372,23 @@ void generateArraySize(MatchExpressionVisitorContext* context,
                 sbe::makeE<sbe::EConstant>(sbe::value::TypeTags::NumberInt64,
                                            sbe::value::bitcastFrom<int64_t>(0))));
 
-        return {std::move(sizeOutput), std::move(traverseStage)};
+        std::vector<EvalExprStagePair> branches;
+
+        // Check that the thing we are about traverse is indeed an array.
+        branches.emplace_back(makeFillEmptyFalse(sbe::makeE<sbe::EFunction>(
+                                  "isArray", sbe::makeEs(sbe::makeE<sbe::EVariable>(inputSlot)))),
+                              EvalStage{});
+
+        branches.emplace_back(std::move(sizeOutput), std::move(traverseStage));
+
+        auto [opOutput, opStage] = generateShortCircuitingLogicalOp(sbe::EPrimBinary::logicAnd,
+                                                                    std::move(branches),
+                                                                    context->planNodeId,
+                                                                    context->slotIdGenerator);
+
+        inputStage = makeLoopJoin(std::move(inputStage), std::move(opStage), context->planNodeId);
+
+        return {std::move(opOutput), std::move(inputStage)};
     };
 
     generatePredicate(context,