SERVER-13039 handle subnodes which require an index during plan enumeration

3 files changed, 190 insertions, 16 deletions

diff --git a/src/mongo/db/query/plan_enumerator.cpp b/src/mongo/db/query/plan_enumerator.cpp
index 3317c516296..99350d8f7d3 100644
--- a/src/mongo/db/query/plan_enumerator.cpp
+++ b/src/mongo/db/query/plan_enumerator.cpp
@@ -36,6 +36,8 @@
 
 namespace {
 
+    using namespace mongo;
+
     std::string getPathPrefix(std::string path) {
         if (mongoutils::str::contains(path, '.')) {
             return mongoutils::str::before(path, '.');
@@ -45,6 +47,15 @@ namespace {
         }
     }
 
+    /**
+     * Returns true if either 'node' or a descendent of 'node'
+     * is a predicate that is required to use an index.
+     */
+    bool isIndexMandatory(const MatchExpression* node) {
+        return CanonicalQuery::countNodes(node, MatchExpression::GEO_NEAR) > 0
+            || CanonicalQuery::countNodes(node, MatchExpression::TEXT) > 0;
+    }
+
 } // namespace
 
 
@@ -256,16 +267,34 @@ namespace mongo {
             IndexToPredMap idxToFirst;
             IndexToPredMap idxToNotFirst;
 
-            // Children that aren't predicates.
+            // Children that aren't predicates, and which do not necessarily need
+            // to use an index.
             vector<MemoID> subnodes;
 
+            // Children that aren't predicates, but which *must* use an index.
+            // (e.g. an OR which contains a TEXT child).
+            vector<MemoID> mandatorySubnodes;
+
             // A list of predicates contained in the subtree rooted at 'node'
             // obtained by traversing deeply through $and and $elemMatch children.
             vector<MatchExpression*> indexedPreds;
 
-            // Partition the childen into the children that aren't predicates
-            // ('subnodes'), and children that are predicates ('indexedPreds').
-            partitionPreds(node, childContext, &indexedPreds, &subnodes);
+            // Partition the childen into the children that aren't predicates which may or may
+            // not be indexed ('subnodes'), children that aren't predicates which must use the
+            // index ('mandatorySubnodes'). and children that are predicates ('indexedPreds').
+            //
+            // We have to get the subnodes with mandatory assignments rather than adding the
+            // mandatory 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.
+            if (!partitionPreds(node, childContext, &indexedPreds,
+                                &subnodes, &mandatorySubnodes)) {
+                return false;
+            }
+
+            if (mandatorySubnodes.size() > 1) {
+                return false;
+            }
 
             // Indices we *must* use.  TEXT or GEO.
             set<IndexID> mandatoryIndices;
@@ -277,8 +306,7 @@ namespace mongo {
 
                 invariant(Indexability::nodeCanUseIndexOnOwnField(child));
 
-                bool childRequiresIndex = (child->matchType() == MatchExpression::GEO_NEAR
-                                        || child->matchType() == MatchExpression::TEXT);
+                bool childRequiresIndex = isIndexMandatory(child);
 
                 RelevantTag* rt = static_cast<RelevantTag*>(child->getTag());
 
@@ -301,7 +329,11 @@ namespace mongo {
             }
 
             // If none of our children can use indices, bail out.
-            if (idxToFirst.empty() && (subnodes.size() == 0)) { return false; }
+            if (idxToFirst.empty()
+                && (subnodes.size() == 0)
+                && (mandatorySubnodes.size() == 0)) {
+                return false;
+            }
 
             // At least one child can use an index, so we can create a memo entry.
             AndAssignment* andAssignment = new AndAssignment();
@@ -312,6 +344,15 @@ namespace mongo {
             // Takes ownership.
             nodeAssignment->andAssignment.reset(andAssignment);
 
+            // Predicates which must use an index might be buried inside
+            // a subnode. Handle that case here.
+            if (1 == mandatorySubnodes.size()) {
+                AndEnumerableState aes;
+                aes.subnodesToIndex.push_back(mandatorySubnodes[0]);
+                andAssignment->choices.push_back(aes);
+                return true;
+            }
+
             // Only near queries and text queries have mandatory indices.
             // In this case there's no point to enumerating anything here; both geoNear
             // and text do fetches internally so we can't use any other indices in conjunction.
@@ -732,10 +773,11 @@ namespace mongo {
         }
     }
 
-    void PlanEnumerator::partitionPreds(MatchExpression* node,
+    bool PlanEnumerator::partitionPreds(MatchExpression* node,
                                         PrepMemoContext context,
                                         vector<MatchExpression*>* indexOut,
-                                        vector<MemoID>* subnodesOut) {
+                                        vector<MemoID>* subnodesOut,
+                                        vector<MemoID>* mandatorySubnodes) {
         for (size_t i = 0; i < node->numChildren(); ++i) {
             MatchExpression* child = node->getChild(i);
             if (Indexability::nodeCanUseIndexOnOwnField(child)) {
@@ -757,17 +799,19 @@ namespace mongo {
                 indexOut->push_back(child);
             }
             else if (Indexability::isBoundsGeneratingNot(child)) {
-                partitionPreds(child, context, indexOut, subnodesOut);
+                partitionPreds(child, context, indexOut, subnodesOut, mandatorySubnodes);
             }
             else if (MatchExpression::ELEM_MATCH_OBJECT == child->matchType()) {
                 PrepMemoContext childContext;
                 childContext.elemMatchExpr = child;
-                partitionPreds(child, childContext, indexOut, subnodesOut);
+                partitionPreds(child, childContext, indexOut, subnodesOut, mandatorySubnodes);
             }
             else if (MatchExpression::AND == child->matchType()) {
-                partitionPreds(child, context, indexOut, subnodesOut);
+                partitionPreds(child, context, indexOut, subnodesOut, mandatorySubnodes);
             }
             else {
+                bool mandatory = isIndexMandatory(child);
+
                 // Recursively prepMemo for the subnode. We fall through
                 // to this case for logical nodes other than AND (e.g. OR)
                 // and for array nodes other than ELEM_MATCH_OBJECT or
@@ -777,10 +821,22 @@ namespace mongo {
                     size_t childID = _nodeToId[child];
 
                     // Output the subnode.
-                    subnodesOut->push_back(childID);
+                    if (mandatory) {
+                        mandatorySubnodes->push_back(childID);
+                    }
+                    else {
+                        subnodesOut->push_back(childID);
+                    }
+                }
+                else if (mandatory) {
+                    // The subnode is mandatory but cannot be indexed. This means
+                    // that the entire AND cannot be indexed either.
+                    return false;
                 }
             }
         }
+
+        return true;
     }
 
     void PlanEnumerator::getMultikeyCompoundablePreds(const MatchExpression* assigned,
diff --git a/src/mongo/db/query/plan_enumerator.h b/src/mongo/db/query/plan_enumerator.h
index 9cd850e2ab0..730d9f2af40 100644
--- a/src/mongo/db/query/plan_enumerator.h
+++ b/src/mongo/db/query/plan_enumerator.h
@@ -251,14 +251,18 @@ namespace mongo {
          * information due to flattening.
          *
          * Nodes that cannot be deeply traversed are returned via the output
-         * vector 'subnodesOut'.
+         * vectors 'subnodesOut' and 'mandatorySubnodes'. Subnodes are "mandatory"
+         * if they *must* use an index (TEXT and GEO).
          *
          * Does not take ownership of arguments.
+         *
+         * Returns false if the AND cannot be indexed. Otherwise returns true.
          */
-        void partitionPreds(MatchExpression* node,
+        bool partitionPreds(MatchExpression* node,
                             PrepMemoContext context,
                             vector<MatchExpression*>* indexOut,
-                            vector<MemoID>* subnodesOut);
+                            vector<MemoID>* subnodesOut,
+                            vector<MemoID>* mandatorySubnodes);
 
         /**
          * Finds a set of predicates that can be safely compounded with 'assigned',
diff --git a/src/mongo/db/query/query_planner_text_test.cpp b/src/mongo/db/query/query_planner_text_test.cpp
index c7d0c5974b1..e6d0a153543 100644
--- a/src/mongo/db/query/query_planner_text_test.cpp
+++ b/src/mongo/db/query/query_planner_text_test.cpp
@@ -439,4 +439,118 @@ namespace {
         assertNumSolutions(1);
     }
 
+    TEST_F(QueryPlannerTest, TextInsideAndWithCompoundIndex) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1 << "_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{$and: [{a: 3}, {$text: {$search: 'foo'}}], a: 3}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{text: {prefix: {a:3}, search: 'foo'}}");
+    }
+
+    // SERVER-13039: Test that we don't generate invalid solutions when the TEXT node
+    // is buried beneath a logical node.
+    TEST_F(QueryPlannerTest, TextInsideOrBasic) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{a: 0, $or: [{_id: 1}, {$text: {$search: 'foo'}}]}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{fetch: {filter: {a:0}, node: {or: {nodes: ["
+                                "{text: {search: 'foo'}}, "
+                                "{ixscan: {filter: null, pattern: {_id: 1}}}]}}}}");
+    }
+
+    // SERVER-13039
+    TEST_F(QueryPlannerTest, TextInsideOrWithAnotherOr) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{$and: [{$or: [{a: 3}, {a: 4}]}, "
+                                  "{$or: [{$text: {$search: 'foo'}}, {a: 5}]}]}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{fetch: {filter: {$or: [{a: 3}, {a: 4}]}, node: "
+                                "{or: {nodes: ["
+                                    "{text: {search: 'foo'}}, "
+                                    "{ixscan: {filter: null, pattern: {a: 1}}}]}}}}");
+    }
+
+    // SERVER-13039
+    TEST_F(QueryPlannerTest, TextInsideOrOfAnd) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{$or: [{a: {$gt: 1, $gt: 2}}, "
+                                 "{a: {$gt: 3}, $text: {$search: 'foo'}}]}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{fetch: {filter: null, node: {or: {nodes: ["
+                                "{ixscan: {filter: null, pattern: {a:1}, bounds: "
+                                    "{a: [[2,Infinity,false,true]]}}}, "
+                                "{fetch: {filter: {a:{$gt:3}}, node: "
+                                    "{text: {search: 'foo'}}}}]}}}}");
+    }
+
+    // SERVER-13039
+    TEST_F(QueryPlannerTest, TextInsideAndOrAnd) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("b" << 1));
+        addIndex(BSON("_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{a: 1, $or: [{a:2}, {b:2}, "
+                                       "{a: 1, $text: {$search: 'foo'}}]}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{fetch: {filter: {a:1}, node: {or: {nodes: ["
+                                "{ixscan: {filter: null, pattern: {a:1}}}, "
+                                "{fetch: {filter: {a:1}, node: {text: {search: 'foo'}}}}, "
+                                "{ixscan: {filter: null, pattern: {b:1}}}]}}}}");
+    }
+
+    // SERVER-13039
+    TEST_F(QueryPlannerTest, TextInsideAndOrAndOr) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{$or: [{a: {$gt: 1, $gt: 2}}, "
+                                 "{a: {$gt: 3}, $or: [{$text: {$search: 'foo'}}, "
+                                                     "{a: 6}]}], "
+                            "a: 5}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{fetch: {filter: {a:5}, node: {or: {nodes: ["
+                                "{ixscan: {filter: null, pattern: {a: 1}}}, "
+                                "{fetch: {filter: {a:{$gt:3}}, node: {or: {nodes: ["
+                                    "{text: {search: 'foo'}}, "
+                                    "{ixscan: {filter: null, pattern: {a: 1}}}]}}}}]}}}}");
+    }
+
+    // If only one branch of the $or can be indexed, then no indexed
+    // solutions are generated, even if one branch is $text.
+    TEST_F(QueryPlannerTest, TextInsideOrOneBranchNotIndexed) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{a: 1, $or: [{b: 2}, {$text: {$search: 'foo'}}]}"));
+
+        assertNumSolutions(0);
+    }
+
+    // If the unindexable $or is not the one containing the $text predicate,
+    // then we should still be able to generate an indexed solution.
+    TEST_F(QueryPlannerTest, TextInsideOrWithAnotherUnindexableOr) {
+        params.options = QueryPlannerParams::NO_TABLE_SCAN;
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("_fts" << "text" << "_ftsx" << 1));
+        runQuery(fromjson("{$and: [{$or: [{a: 1}, {b: 1}]}, "
+                                  "{$or: [{a: 2}, {$text: {$search: 'foo'}}]}]}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{fetch: {filter: {$or:[{a:1},{b:1}]}, node: {or: {nodes: ["
+                                "{text: {search: 'foo'}}, "
+                                "{ixscan: {filter: null, pattern: {a:1}}}]}}}}");
+    }
+
 }  // namespace