SERVER-62058 Implement $_internalExpr in SBE

src/mongo/db/matcher/expression_internal_expr_comparison.h

@@ -68,14 +68,21 @@ namespace mongo {
 template <typename T>
 class InternalExprComparisonMatchExpression : public ComparisonMatchExpressionBase {
 public:
+    /**
+     * Constructor. 'mustExecute' indicates whether this match expression is forced to be evaluated.
+     * This defaults to false since InternalExpr* expressions are permitted to return false
+     * positives and are sometimes elided from execution plans.
+     */
     InternalExprComparisonMatchExpression(MatchType type,
                                           boost::optional<StringData> path,
-                                          BSONElement value)
+                                          BSONElement value,
+                                          bool mustExecute)
         : ComparisonMatchExpressionBase(type,
                                         path,
                                         Value(value),
                                         ElementPath::LeafArrayBehavior::kNoTraversal,
-                                        ElementPath::NonLeafArrayBehavior::kMatchSubpath) {
+                                        ElementPath::NonLeafArrayBehavior::kMatchSubpath),
+          _mustExecute(mustExecute) {
         invariant(_rhs.type() != BSONType::Undefined);
         invariant(_rhs.type() != BSONType::Array);
     }
@@ -124,6 +131,19 @@ class InternalExprComparisonMatchExpression : public ComparisonMatchExpressionBa
     StringData name() const final {
         return T::kName;
     };
+
+    /*
+     * Return whether this node is marked as 'mustExecute'. InternalExpr* nodes without this flag
+     * set may get elided from the execution plan, since InternalExpr* expressions are permitted to
+     * return false positives. Setting the flag forces the execution system to evaluate these
+     * predicates.
+     */
+    bool mustExecute() const {
+        return _mustExecute;
+    }
+
+private:
+    bool _mustExecute;
 };
 
 
@@ -132,9 +152,11 @@ class InternalExprEqMatchExpression final
 public:
     static constexpr StringData kName = "$_internalExprEq"_sd;
 
-    InternalExprEqMatchExpression(boost::optional<StringData> path, BSONElement value)
+    InternalExprEqMatchExpression(boost::optional<StringData> path,
+                                  BSONElement value,
+                                  bool mustExecute = false)
         : InternalExprComparisonMatchExpression<InternalExprEqMatchExpression>(
-              MatchType::INTERNAL_EXPR_EQ, path, value) {}
+              MatchType::INTERNAL_EXPR_EQ, path, value, mustExecute) {}
 
     void acceptVisitor(MatchExpressionMutableVisitor* visitor) final {
         visitor->visit(this);
@@ -150,9 +172,11 @@ class InternalExprGTMatchExpression final
 public:
     static constexpr StringData kName = "$_internalExprGt"_sd;
 
-    InternalExprGTMatchExpression(boost::optional<StringData> path, BSONElement value)
+    InternalExprGTMatchExpression(boost::optional<StringData> path,
+                                  BSONElement value,
+                                  bool mustExecute = false)
         : InternalExprComparisonMatchExpression<InternalExprGTMatchExpression>(
-              MatchType::INTERNAL_EXPR_GT, path, value) {}
+              MatchType::INTERNAL_EXPR_GT, path, value, mustExecute) {}
 
 
     void acceptVisitor(MatchExpressionMutableVisitor* visitor) final {
@@ -169,9 +193,11 @@ class InternalExprGTEMatchExpression final
 public:
     static constexpr StringData kName = "$_internalExprGte"_sd;
 
-    InternalExprGTEMatchExpression(boost::optional<StringData> path, BSONElement value)
+    InternalExprGTEMatchExpression(boost::optional<StringData> path,
+                                   BSONElement value,
+                                   bool mustExecute = false)
         : InternalExprComparisonMatchExpression<InternalExprGTEMatchExpression>(
-              MatchType::INTERNAL_EXPR_GTE, path, value) {}
+              MatchType::INTERNAL_EXPR_GTE, path, value, mustExecute) {}
 
     void acceptVisitor(MatchExpressionMutableVisitor* visitor) final {
         visitor->visit(this);
@@ -187,9 +213,11 @@ class InternalExprLTMatchExpression final
 public:
     static constexpr StringData kName = "$_internalExprLt"_sd;
 
-    InternalExprLTMatchExpression(boost::optional<StringData> path, BSONElement value)
+    InternalExprLTMatchExpression(boost::optional<StringData> path,
+                                  BSONElement value,
+                                  bool mustExecute = false)
         : InternalExprComparisonMatchExpression<InternalExprLTMatchExpression>(
-              MatchType::INTERNAL_EXPR_LT, path, value) {}
+              MatchType::INTERNAL_EXPR_LT, path, value, mustExecute) {}
 
     void acceptVisitor(MatchExpressionMutableVisitor* visitor) final {
         visitor->visit(this);
@@ -205,9 +233,11 @@ class InternalExprLTEMatchExpression final
 public:
     static constexpr StringData kName = "$_internalExprLte"_sd;
 
-    InternalExprLTEMatchExpression(boost::optional<StringData> path, BSONElement value)
+    InternalExprLTEMatchExpression(boost::optional<StringData> path,
+                                   BSONElement value,
+                                   bool mustExecute = false)
         : InternalExprComparisonMatchExpression<InternalExprLTEMatchExpression>(
-              MatchType::INTERNAL_EXPR_LTE, path, value) {}
+              MatchType::INTERNAL_EXPR_LTE, path, value, mustExecute) {}
 
     void acceptVisitor(MatchExpressionMutableVisitor* visitor) final {
         visitor->visit(this);

src/mongo/db/matcher/expression_leaf.h

@@ -69,30 +69,29 @@ namespace mongo {
 class CollatorInterface;
 
 /**
- *  A struct primarily used to make input parameters for makePredicate() function. The
- * 'MatchExprType' helps in determining which MatchExpression to make in makePredicate() function.
+ * Makes a conjunction of the given predicates
  */
-template <typename MatchExprType, typename ValueType = BSONElement>
-struct MatchExprPredicate {
-    MatchExprPredicate(StringData path_, ValueType value_) : path(path_), value(value_){};
-    StringData path;
-    ValueType value;
-};
+template <typename... Ts>
+inline auto makeAnd(Ts&&... pack) {
+    auto predicates = makeVector<std::unique_ptr<MatchExpression>>(std::forward<Ts>(pack)...);
+    return std::make_unique<AndMatchExpression>(std::move(predicates));
+}
 
 /**
- * Helper function to make $and predicates based on the set of predicates passed as parameters.
+ * Makes a disjunction of the given predicates.
+ *
+ * - The result is non-null; it may be an OrMatchExpression with zero children.
+ * - Any trivially-false arguments are omitted.
+ * - If only one argument is nontrivial, returns that argument rather than adding an extra
+ *   OrMatchExpression around it.
  */
-template <typename T, typename ValueType, typename... Targs, typename... ValueTypeArgs>
-std::unique_ptr<MatchExpression> makePredicate(
-    MatchExprPredicate<T, ValueType> predicate,
-    MatchExprPredicate<Targs, ValueTypeArgs>... predicates) {
-    if constexpr (sizeof...(predicates) > 0) {
-        return std::make_unique<AndMatchExpression>(makeVector<std::unique_ptr<MatchExpression>>(
-            std::make_unique<T>(predicate.path, predicate.value),
-            (std::make_unique<Targs>(predicates.path, predicates.value))...));
-    } else {
-        return std::make_unique<T>(predicate.path, predicate.value);
-    }
+template <typename... Ts>
+inline auto makeOr(Ts&&... pack) {
+    auto predicates = makeVector<std::unique_ptr<MatchExpression>>(std::forward<Ts>(pack)...);
+    auto newEnd = std::remove_if(
+        predicates.begin(), predicates.end(), [](auto& node) { return node->isTriviallyFalse(); });
+    predicates.erase(newEnd, predicates.end());
+    return std::make_unique<OrMatchExpression>(std::move(predicates));
 }
 
 class LeafMatchExpression : public PathMatchExpression {

src/mongo/db/query/sbe_expression_bm.cpp

@@ -136,6 +136,7 @@ class SbeExpressionBenchmarkFixture : public ExpressionBenchmarkFixture {
             &_spoolIdGenerator,
             &_inListsSet,
             &_collatorsMap,
+            _expCtx,
             false /* needsMerge */,
             false /* allowDiskUse */
         };
@@ -199,6 +200,7 @@ class SbeExpressionBenchmarkFixture : public ExpressionBenchmarkFixture {
     sbe::value::SpoolIdGenerator _spoolIdGenerator;
     stage_builder::StageBuilderState::InListsSet _inListsSet;
     stage_builder::StageBuilderState::CollatorsMap _collatorsMap;
+    boost::intrusive_ptr<ExpressionContext> _expCtx;
 
     sbe::value::SlotId _inputSlotId;
     std::unique_ptr<TimeZoneDatabase> _timeZoneDB;

src/mongo/db/query/sbe_stage_builder.cpp

@@ -461,6 +461,7 @@ SlotBasedStageBuilder::SlotBasedStageBuilder(OperationContext* opCtx,
              &_spoolIdGenerator,
              &_inListsSet,
              &_collatorMap,
+             _cq.getExpCtx(),
              _cq.getExpCtx()->needsMerge,
              _cq.getExpCtx()->allowDiskUse) {
     // Initialize '_data->queryCollator'.

src/mongo/db/query/sbe_stage_builder_filter.cpp

@@ -956,32 +956,84 @@ class MatchExpressionPostVisitor final : public MatchExpressionConstVisitor {
         generatePredicate(_context, *expr->fieldRef(), makePredicate, traversalMode, hasNull);
     }
 
-    // The following are no-ops. The internal expr comparison match expression are produced
-    // internally by rewriting an $expr expression to an AND($expr, $_internalExpr[OP]), which can
-    // later be eliminated by via a conversion into EXACT index bounds, or remains present. In the
-    // latter case we can simply ignore it, as the result of AND($expr, $_internalExpr[OP]) is equal
-    // to just $expr.
-    //
-    // TODO SERVER-62058: This is a correct implementation for the time-series loose filter since
-    // the event filter should be able to filter out non-matching results but will not be
-    // performant. But this is an incorrect implementation for the time-series tight filter since
-    // if the tight filter evalutes to 'true', then the event filter will not be evaluated at all.
-    // Implement the exact behavior for the better performance of the loose filter and the
-    // correctness of the tight filter.
+    void translateExprComparison(const ComparisonMatchExpressionBase* expr, bool mustExecute) {
+        /**
+         * Since InternalExpr* are permitted to return false positives, we simply compile this to an
+         * always "true" expression. In practice, most of the time when an InternalExpr is not
+         * marked with 'mustExecute' it is accompanied by a more precise check which removes the
+         * false positives.
+         */
+        if (!mustExecute) {
+            generateAlwaysBoolean(_context, true);
+            return;
+        }
+
+        ExpressionCompare::CmpOp cmp = [&]() {
+            switch (expr->matchType()) {
+                case MatchExpression::MatchType::INTERNAL_EXPR_EQ:
+                    return ExpressionCompare::CmpOp::EQ;
+                case MatchExpression::MatchType::INTERNAL_EXPR_GT:
+                    return ExpressionCompare::CmpOp::GT;
+                case MatchExpression::MatchType::INTERNAL_EXPR_GTE:
+                    return ExpressionCompare::CmpOp::GTE;
+                case MatchExpression::MatchType::INTERNAL_EXPR_LT:
+                    return ExpressionCompare::CmpOp::LT;
+                case MatchExpression::MatchType::INTERNAL_EXPR_LTE:
+                    return ExpressionCompare::CmpOp::LTE;
+                default:
+                    // Only $expr expressions supported.
+                    MONGO_UNREACHABLE_TASSERT(6205800);
+            }
+        }();
+
+        auto expCtx = _context->state.expCtx;
+        invariant(expCtx);
+
+        auto fieldPathExpr = ExpressionFieldPath::createPathFromString(
+            expCtx.get(), expr->fieldRef()->dottedField().toString(), expCtx->variablesParseState);
+        auto translatedFieldPathExpr = generateExpression(
+            _context->state, fieldPathExpr.get(), _context->rootSlot, _context->slots);
+
+        SbExprBuilder b(_context->state);
+        auto frameId = _context->state.frameIdGenerator->generate();
+        auto newVar = b.makeVariable(frameId, 0);
+
+        auto cmpExpr = ExpressionCompare::create(
+            expCtx.get(),
+            cmp,
+            fieldPathExpr,
+            ExpressionConstant::create(expCtx.get(), Value(expr->getData())));
+
+        auto translatedCmpExpr =
+            generateExpression(_context->state, cmpExpr.get(), _context->rootSlot, _context->slots);
+
+        auto isArrayExpr = b.makeIf(b.makeFillEmptyTrue(b.makeFunction("isArray", newVar.clone())),
+                                    b.makeBoolConstant(true),
+                                    std::move(translatedCmpExpr));
+
+        auto cmpWArrayCheckExpr = b.makeLet(
+            frameId, SbExpr::makeSeq(std::move(translatedFieldPathExpr)), std::move(isArrayExpr));
+
+        auto& frame = _context->topFrame();
+        // Convert the result of the '{$expr: ..}' expression to a boolean value.
+        frame.pushExpr(
+            b.makeFillEmptyFalse(b.makeFunction("coerceToBool"_sd, std::move(cmpWArrayCheckExpr))));
+    }
+
     void visit(const InternalExprEqMatchExpression* expr) final {
-        generateAlwaysBoolean(_context, true);
+        translateExprComparison(expr, expr->mustExecute());
     }
     void visit(const InternalExprGTMatchExpression* expr) final {
-        generateAlwaysBoolean(_context, true);
+        translateExprComparison(expr, expr->mustExecute());
     }
     void visit(const InternalExprGTEMatchExpression* expr) final {
-        generateAlwaysBoolean(_context, true);
+        translateExprComparison(expr, expr->mustExecute());
     }
     void visit(const InternalExprLTMatchExpression* expr) final {
-        generateAlwaysBoolean(_context, true);
+        translateExprComparison(expr, expr->mustExecute());
     }
     void visit(const InternalExprLTEMatchExpression* expr) final {
-        generateAlwaysBoolean(_context, true);
+        translateExprComparison(expr, expr->mustExecute());
     }
 
     void visit(const InternalEqHashedKey* expr) final {}

src/mongo/db/query/sbe_stage_builder_state.h

@@ -60,6 +60,7 @@ struct StageBuilderState {
                       sbe::value::SpoolIdGenerator* spoolIdGenerator,
                       InListsSet* inListsSet,
                       CollatorsMap* collatorsMap,
+                      boost::intrusive_ptr<ExpressionContext> expCtx,
                       bool needsMerge,
                       bool allowDiskUse)
         : slotIdGenerator{slotIdGenerator},
@@ -71,6 +72,7 @@ struct StageBuilderState {
           env{env},
           data{data},
           variables{variables},
+          expCtx{expCtx},
           needsMerge{needsMerge},
           allowDiskUse{allowDiskUse} {}
 
@@ -129,6 +131,9 @@ struct StageBuilderState {
     PlanStageStaticData* const data;
 
     const Variables& variables;
+
+    boost::intrusive_ptr<ExpressionContext> expCtx;
+
     // When the mongos splits $group stage and sends it to shards, it adds 'needsMerge'/'fromMongs'
     // flags to true so that shards can sends special partial aggregation results to the mongos.
     bool needsMerge;