SERVER-51592 Added support for anyElementTrue in SBE

5 files changed, 69 insertions, 9 deletions

diff --git a/jstests/aggregation/expressions/any_element_true.js b/jstests/aggregation/expressions/any_element_true.js
index cd9217c7700..05c40d0ef39 100644
--- a/jstests/aggregation/expressions/any_element_true.js
+++ b/jstests/aggregation/expressions/any_element_true.js
@@ -4,8 +4,18 @@
 (function() {
 const coll = db.any_element_true;
 coll.drop();
-assert.commandWorked(
-    coll.insert({_id: 0, allTrue: [true, true], someTrue: [true, false], noneTrue: [0, false]}));
+assert.commandWorked(coll.insert({
+    _id: 0,
+    allTrue: [true, true],
+    someTrue: [true, false],
+    noneTrue: [0, false],
+    nonArray: 1,
+    nullInput: [null],
+    undefinedInput: [undefined],
+    undefinedTrue: [undefined, true],
+    nullTrue: [null, true],
+    empty: []
+}));
 
 function testOp(expression, expected) {
     const results = coll.aggregate([{$project: {_id: 0, result: expression}}]).toArray();
@@ -15,6 +25,12 @@ function testOp(expression, expected) {
     assert.eq(loneResult.result, expected, loneResult);
 }
 
+function assertThrows(expression) {
+    const error =
+        assert.throws(() => coll.aggregate([{$project: {_id: 0, result: expression}}]).toArray());
+    assert.commandFailedWithCode(error, 5159200);
+}
+
 testOp({$anyElementTrue: {$literal: [true, true]}}, true);
 testOp({$anyElementTrue: {$literal: [1, true]}}, true);
 testOp({$anyElementTrue: {$literal: [true, 0]}}, true);
@@ -23,4 +39,14 @@ testOp({$anyElementTrue: {$literal: [false, 0, false]}}, false);
 testOp({$anyElementTrue: "$allTrue"}, true);
 testOp({$anyElementTrue: "$someTrue"}, true);
 testOp({$anyElementTrue: "$noneTrue"}, false);
+testOp({$anyElementTrue: ["$noneTrue"]}, false);
+testOp({$anyElementTrue: [["$non_existent_field"]]}, false);
+testOp({$anyElementTrue: [["$non_existent_field", true]]}, true);
+testOp({$anyElementTrue: "$nullInput"}, false);
+testOp({$anyElementTrue: "$undefinedInput"}, false);
+testOp({$anyElementTrue: "$undefinedTrue"}, true);
+testOp({$anyElementTrue: "$nullTrue"}, true);
+testOp({$anyElementTrue: "$empty"}, false);
+assertThrows({$anyElementTrue: "$nonArray"});
+assertThrows({$anyElementTrue: ["$non_existent_field"]});
 }());
diff --git a/jstests/libs/sbe_assert_error_override.js b/jstests/libs/sbe_assert_error_override.js
index d8df0822788..638d37858df 100644
--- a/jstests/libs/sbe_assert_error_override.js
+++ b/jstests/libs/sbe_assert_error_override.js
@@ -113,6 +113,7 @@ const equivalentErrorCodesList = [
     [5338802, 5439016],
     [5687301, 5687400],
     [5687302, 5687401],
+    [17041, 5159200],
 ];
 
 // This map is generated based on the contents of 'equivalentErrorCodesList'. This map should _not_
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index f696152ec59..40bc488771d 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -1001,13 +1001,9 @@ public:
 class ExpressionAnyElementTrue final : public ExpressionFixedArity<ExpressionAnyElementTrue, 1> {
 public:
     explicit ExpressionAnyElementTrue(ExpressionContext* const expCtx)
-        : ExpressionFixedArity<ExpressionAnyElementTrue, 1>(expCtx) {
-        expCtx->sbeCompatible = false;
-    }
+        : ExpressionFixedArity<ExpressionAnyElementTrue, 1>(expCtx) {}
     ExpressionAnyElementTrue(ExpressionContext* const expCtx, ExpressionVector&& children)
-        : ExpressionFixedArity<ExpressionAnyElementTrue, 1>(expCtx, std::move(children)) {
-        expCtx->sbeCompatible = false;
-    }
+        : ExpressionFixedArity<ExpressionAnyElementTrue, 1>(expCtx, std::move(children)) {}
 
     Value evaluate(const Document& root, Variables* variables) const final;
     const char* getOpName() const final;
diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp
index 4e25a8e18a9..90dee12759a 100644
--- a/src/mongo/db/query/sbe_stage_builder_expression.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp
@@ -849,7 +849,40 @@ public:
         visitMultiBranchLogicExpression(expr, sbe::EPrimBinary::logicAnd);
     }
     void visit(const ExpressionAnyElementTrue* expr) final {
-        unsupportedExpression(expr->getOpName());
+        auto [inputSlot, stage] = projectEvalExpr(_context->popEvalExpr(),
+                                                  _context->extractCurrentEvalStage(),
+                                                  _context->planNodeId,
+                                                  _context->state.slotIdGenerator);
+
+        auto fromBranch = makeFilter<false>(
+            std::move(stage),
+            makeBinaryOp(sbe::EPrimBinary::logicOr,
+                         makeFillEmptyFalse(makeFunction("isArray", makeVariable(inputSlot))),
+                         makeFail(5159200, "$anyElementTrue's argument must be an array")),
+            _context->planNodeId);
+
+        auto innerOutputSlot = _context->state.slotId();
+        auto innerBranch = makeProject(makeLimitCoScanStage(_context->planNodeId),
+                                       _context->planNodeId,
+                                       innerOutputSlot,
+                                       generateCoerceToBoolExpression(inputSlot));
+
+        auto traverseSlot = _context->state.slotId();
+        auto traverseStage = makeTraverse(std::move(fromBranch),
+                                          std::move(innerBranch),
+                                          inputSlot,
+                                          traverseSlot,
+                                          innerOutputSlot,
+                                          makeBinaryOp(sbe::EPrimBinary::logicOr,
+                                                       makeVariable(traverseSlot),
+                                                       makeVariable(innerOutputSlot)),
+                                          makeVariable(traverseSlot),
+                                          _context->planNodeId,
+                                          1,
+                                          _context->getLexicalEnvironment());
+
+        _context->pushExpr(makeFillEmptyFalse(makeVariable(traverseSlot)),
+                           std::move(traverseStage));
     }
     void visit(const ExpressionArray* expr) final {
         unsupportedExpression(expr->getOpName());
diff --git a/src/mongo/db/query/sbe_stage_builder_helpers.h b/src/mongo/db/query/sbe_stage_builder_helpers.h
index c5a63a3b002..46f41471ac0 100644
--- a/src/mongo/db/query/sbe_stage_builder_helpers.h
+++ b/src/mongo/db/query/sbe_stage_builder_helpers.h
@@ -216,6 +216,10 @@ std::unique_ptr<sbe::EExpression> makeVariable(sbe::value::SlotId slotId);
 
 std::unique_ptr<sbe::EExpression> makeVariable(sbe::FrameId frameId, sbe::value::SlotId slotId);
 
+inline auto makeFail(int code, StringData errorMessage) {
+    return sbe::makeE<sbe::EFail>(ErrorCodes::Error{code}, errorMessage);
+}
+
 /**
  * Check if expression returns Nothing and return null if so. Otherwise, return the expression.
  */