SERVER-70597 Don't generate OrPushdown tags when optimization is disabled

3 files changed, 47 insertions, 6 deletions

diff --git a/src/mongo/db/query/plan_enumerator.cpp b/src/mongo/db/query/plan_enumerator.cpp
index 0eedf0e07b2..26dad772006 100644
--- a/src/mongo/db/query/plan_enumerator.cpp
+++ b/src/mongo/db/query/plan_enumerator.cpp
@@ -263,7 +263,8 @@ PlanEnumerator::PlanEnumerator(const PlanEnumeratorParams& params)
       _ixisect(params.intersect),
       _enumerateOrChildrenLockstep(params.enumerateOrChildrenLockstep),
       _orLimit(params.maxSolutionsPerOr),
-      _intersectLimit(params.maxIntersectPerAnd) {}
+      _intersectLimit(params.maxIntersectPerAnd),
+      _disableOrPushdown(params.disableOrPushdown) {}
 
 PlanEnumerator::~PlanEnumerator() {
     typedef stdx::unordered_map<MemoID, NodeAssignment*> MemoMap;
@@ -528,10 +529,14 @@ bool PlanEnumerator::prepMemo(MatchExpression* node, PrepMemoContext context) {
         // preds to 'indexedPreds'. Adding the mandatory preds directly to 'indexedPreds' would lead
         // to problems such as pulling a predicate beneath an OR into a set joined by an AND.
         getIndexedPreds(node, childContext, &indexedPreds);
-        // Pass in the indexed predicates as outside predicates when prepping the subnodes.
+        // Pass in the indexed predicates as outside predicates when prepping the subnodes. But if
+        // match expression optimization is disabled, skip this part: we don't want to do
+        // OR-pushdown because it relies on the expression being canonicalized.
         auto childContextCopy = childContext;
-        for (auto pred : indexedPreds) {
-            childContextCopy.outsidePreds[pred] = OutsidePredRoute{};
+        if (MONGO_likely(!_disableOrPushdown)) {
+            for (auto pred : indexedPreds) {
+                childContextCopy.outsidePreds[pred] = OutsidePredRoute{};
+            }
         }
         if (!prepSubNodes(node, childContextCopy, &subnodes, &mandatorySubnodes)) {
             return false;
@@ -835,6 +840,13 @@ void PlanEnumerator::assignPredicate(
     MatchExpression* pred,
     size_t position,
     OneIndexAssignment* indexAssignment) {
+    if (MONGO_unlikely(_disableOrPushdown)) {
+        // If match expression optimization is disabled, we also disable OR-pushdown,
+        // so we should never get 'outsidePreds' here.
+        tassert(7059700,
+                "Tried to do OR-pushdown despite disableMatchExpressionOptimization",
+                outsidePreds.empty());
+    }
     if (outsidePreds.find(pred) != outsidePreds.end()) {
         OrPushdownTag::Destination dest;
         dest.route = outsidePreds.at(pred).route;
diff --git a/src/mongo/db/query/plan_enumerator.h b/src/mongo/db/query/plan_enumerator.h
index 60344f0c9ee..36a01f52d50 100644
--- a/src/mongo/db/query/plan_enumerator.h
+++ b/src/mongo/db/query/plan_enumerator.h
@@ -44,7 +44,8 @@ namespace mongo {
 struct PlanEnumeratorParams {
     PlanEnumeratorParams()
         : maxSolutionsPerOr(internalQueryEnumerationMaxOrSolutions.load()),
-          maxIntersectPerAnd(internalQueryEnumerationMaxIntersectPerAnd.load()) {}
+          maxIntersectPerAnd(internalQueryEnumerationMaxIntersectPerAnd.load()),
+          disableOrPushdown(disableMatchExpressionOptimization.shouldFail()) {}
 
     // Do we provide solutions that use more indices than the minimum required to provide
     // an indexed solution?
@@ -69,6 +70,11 @@ struct PlanEnumeratorParams {
     // all-pairs approach, we could wind up creating a lot of enumeration possibilities for
     // certain inputs.
     size_t maxIntersectPerAnd;
+
+    // Whether to disable OR-pushdown optimization. OR-pushdown assumes that the expression has been
+    // simplified: for example, that single-child $or nodes are unwrapped. To avoid this, when
+    // the 'disableMatchExpressionOptimization' failpoint is set, we also disable OR-pushdown.
+    bool disableOrPushdown;
 };
 
 /**
@@ -594,6 +600,9 @@ private:
 
     // How many things do we want from each AND?
     size_t _intersectLimit;
+
+    // Whether we should disable OR-pushdown optimization.
+    const bool _disableOrPushdown;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_planner_tree_test.cpp b/src/mongo/db/query/query_planner_tree_test.cpp
index a7e08ff6488..6708afda6a4 100644
--- a/src/mongo/db/query/query_planner_tree_test.cpp
+++ b/src/mongo/db/query/query_planner_tree_test.cpp
@@ -433,7 +433,7 @@ TEST_F(QueryPlannerTest, RootedOrOfAndDontCollapseDifferentBounds) {
         "bounds: {c: [[3,3,true,true]], d: [[4,4,true,true]]}}}]}}}}");
 }
 
-TEST_F(QueryPlannerTest, DontCrashTryingToPushToSingleChildIndexedOr) {
+TEST_F(QueryPlannerTest, DontCrashTryingToPushToSingleChildIndexedOr1) {
     FailPointEnableBlock failPoint("disableMatchExpressionOptimization");
     addIndex(BSON("indexed" << 1));
     runQuery(
@@ -446,6 +446,26 @@ TEST_F(QueryPlannerTest, DontCrashTryingToPushToSingleChildIndexedOr) {
     assertNumSolutions(3U);
 }
 
+TEST_F(QueryPlannerTest, DontCrashTryingToPushToSingleChildIndexedOr2) {
+    // Test that queries with single-child $and, $or do not crash when match-expression optimization
+    // is disabled. Normally these single-child nodes are eliminated, so when they are left in place
+    // it can confuse OR-pushdown optimization.
+    //
+    // Originally designed to reproduce SERVER-70597, which would only happen when the
+    // INDEX_INTERSECTION option is enabled.
+    FailPointEnableBlock failPoint("disableMatchExpressionOptimization");
+    addIndex(BSON("a" << 1 << "b" << 1));
+
+    params.options |= QueryPlannerParams::INDEX_INTERSECTION;
+    runQuery(
+        fromjson("{ $and : [\n"
+                 "      { $and : [ { a : 2 } ] },\n"
+                 "      { $or : [ { b : 3 } ] }\n"
+                 "  ] }"));
+
+    assertNumSolutions(2U);
+}
+
 // SERVER-13960: properly handle $or with a mix of exact and inexact predicates.
 TEST_F(QueryPlannerTest, OrInexactWithExact) {
     addIndex(BSON("name" << 1));