SERVER-13789 recursively build data access for subnodes beneath elemMatch

(cherry picked from commit eb97470c0507020e3c0f1cc47751feaa777b5d2e)

3 files changed, 131 insertions, 9 deletions

diff --git a/src/mongo/db/query/planner_access.cpp b/src/mongo/db/query/planner_access.cpp
index fc0fb1c2030..413860ac396 100644
--- a/src/mongo/db/query/planner_access.cpp
+++ b/src/mongo/db/query/planner_access.cpp
@@ -482,16 +482,20 @@ namespace mongo {
 
     // static
     void QueryPlannerAccess::findElemMatchChildren(const MatchExpression* node,
-                                                   vector<MatchExpression*>* out) {
+                                                   vector<MatchExpression*>* out,
+                                                   vector<MatchExpression*>* subnodesOut) {
         for (size_t i = 0; i < node->numChildren(); ++i) {
             MatchExpression* child = node->getChild(i);
-            if (Indexability::nodeCanUseIndexOnOwnField(child) &&
+            if (Indexability::isBoundsGenerating(child) &&
                 NULL != child->getTag()) {
                 out->push_back(child);
             }
             else if (MatchExpression::AND == child->matchType() ||
                      Indexability::arrayUsesIndexOnChildren(child)) {
-                findElemMatchChildren(child, out);
+                findElemMatchChildren(child, out, subnodesOut);
+            }
+            else if (NULL != child->getTag()) {
+                subnodesOut->push_back(child);
             }
         }
     }
@@ -545,10 +549,42 @@ namespace mongo {
                     // predicates from inside the $elemMatch, and try to merge them with
                     // the current index scan.
 
-                    // Populate 'emChildren' with tagged predicates from inside the
-                    // tree rooted at 'child.
+                    // Contains tagged predicates from inside the tree rooted at 'child'
+                    // which are logically part of the AND.
                     vector<MatchExpression*> emChildren;
-                    findElemMatchChildren(child, &emChildren);
+
+                    // Contains tagged nodes that are not logically part of the AND and
+                    // cannot use the index directly (e.g. OR nodes which are tagged to
+                    // be indexed).
+                    vector<MatchExpression*> emSubnodes;
+
+                    // Populate 'emChildren' and 'emSubnodes'.
+                    findElemMatchChildren(child, &emChildren, &emSubnodes);
+
+                    // Recursively build data access for the nodes inside 'emSubnodes'.
+                    for (size_t i = 0; i < emSubnodes.size(); ++i) {
+                        MatchExpression* subnode = emSubnodes[i];
+
+                        if (!Indexability::isBoundsGenerating(subnode)) {
+                            // Must pass true for 'inArrayOperator' because the subnode is
+                            // beneath an ELEM_MATCH_OBJECT.
+                            QuerySolutionNode* childSolution = buildIndexedDataAccess(query,
+                                                                                      subnode,
+                                                                                      true,
+                                                                                      indices);
+
+                            // buildIndexedDataAccess(...) returns NULL in error conditions, when
+                            // it is unable to construct a query solution from a tagged match
+                            // expression tree. If we are unable to construct a solution according
+                            // to the instructions from the enumerator, then we  bail out early
+                            // (by returning false) rather than continuing on and potentially
+                            // constructing an invalid solution tree.
+                            if (NULL == childSolution) { return false; }
+
+                            // Output the resulting solution tree.
+                            out->push_back(childSolution);
+                        }
+                    }
 
                     // For each predicate in 'emChildren', try to merge it with the
                     // current index scan.
diff --git a/src/mongo/db/query/planner_access.h b/src/mongo/db/query/planner_access.h
index f2d002ae370..71dd60d9f39 100644
--- a/src/mongo/db/query/planner_access.h
+++ b/src/mongo/db/query/planner_access.h
@@ -163,11 +163,15 @@ namespace mongo {
          * finding all predicates that can use an index directly and returning
          * them in the out-parameter vector 'out'.
          *
-         * Traverses only through $and and $elemMatch nodes, not through other
-         * logical or array nodes like $or and $all.
+         * Traverses only through $and and array nodes like $all.
+         *
+         * Other nodes (i.e. nodes which cannot use an index directly, and which are
+         * neither $and nor array nodes) are returned in 'subnodesOut' if they are
+         * tagged to use an index.
          */
         static void findElemMatchChildren(const MatchExpression* node,
-                                          vector<MatchExpression*>* out);
+                                          vector<MatchExpression*>* out,
+                                          vector<MatchExpression*>* subnodesOut);
 
         /**
          * Helper used by buildIndexedAnd and buildIndexedOr.
diff --git a/src/mongo/db/query/query_planner_test.cpp b/src/mongo/db/query/query_planner_test.cpp
index d56efe11051..871f0c73c94 100644
--- a/src/mongo/db/query/query_planner_test.cpp
+++ b/src/mongo/db/query/query_planner_test.cpp
@@ -1416,6 +1416,88 @@ namespace {
         runQuery(fromjson("{a: {$elemMatch: {$not: {$gte: 6}}}}"));
     }
 
+    // SERVER-13789
+    TEST_F(QueryPlannerTest, ElemMatchIndexedNestedOr) {
+        addIndex(BSON("bar.baz" << 1));
+        runQuery(fromjson("{foo: 1, $and: [{bar: {$elemMatch: {$or: [{baz: 2}]}}}]}"));
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{cscan: {dir: 1}}");
+        assertSolutionExists("{fetch: {filter: {$and: [{foo:1},"
+                                               "{bar:{$elemMatch:{$or:[{baz:2}]}}}]}, "
+                                "node: {ixscan: {pattern: {'bar.baz': 1}, "
+                                                "bounds: {'bar.baz': [[2,2,true,true]]}}}}}");
+    }
+
+    // SERVER-13789
+    TEST_F(QueryPlannerTest, ElemMatchIndexedNestedOrMultiplePreds) {
+        addIndex(BSON("bar.baz" << 1));
+        addIndex(BSON("bar.z" << 1));
+        runQuery(fromjson("{foo: 1, $and: [{bar: {$elemMatch: {$or: [{baz: 2}, {z: 3}]}}}]}"));
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{cscan: {dir: 1}}");
+        assertSolutionExists("{fetch: {filter: {$and: [{foo:1},"
+                                               "{bar:{$elemMatch:{$or:[{baz:2},{z:3}]}}}]}, "
+                                "node: {or: {nodes: ["
+                                    "{ixscan: {pattern: {'bar.baz': 1}, "
+                                              "bounds: {'bar.baz': [[2,2,true,true]]}}},"
+                                    "{ixscan: {pattern: {'bar.z': 1}, "
+                                              "bounds: {'bar.z': [[3,3,true,true]]}}}]}}}}");
+    }
+
+    // SERVER-13789: Ensure that we properly compound in the multikey case when an
+    // $or is beneath an $elemMatch.
+    TEST_F(QueryPlannerTest, ElemMatchIndexedNestedOrMultikey) {
+        // true means multikey
+        addIndex(BSON("bar.baz" << 1 << "bar.z" << 1), true);
+        runQuery(fromjson("{foo: 1, $and: [{bar: {$elemMatch: {$or: [{baz: 2, z: 3}]}}}]}"));
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{cscan: {dir: 1}}");
+        assertSolutionExists("{fetch: {filter: {$and: [{foo:1},"
+                                "{bar: {$elemMatch: {$or: [{$and: [{baz:2}, {z:3}]}]}}}]},"
+                                "node: {ixscan: {pattern: {'bar.baz': 1, 'bar.z': 1}, "
+                                                "bounds: {'bar.baz': [[2,2,true,true]],"
+                                                         "'bar.z': [[3,3,true,true]]}}}}}");
+    }
+
+    // SERVER-13789: Right now we don't index $nor, but make sure that the planner
+    // doesn't get confused by a $nor beneath an $elemMatch.
+    TEST_F(QueryPlannerTest, ElemMatchIndexedNestedNor) {
+        addIndex(BSON("bar.baz" << 1));
+        runQuery(fromjson("{foo: 1, $and: [{bar: {$elemMatch: {$nor: [{baz: 2}, {baz: 3}]}}}]}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{cscan: {dir: 1}}");
+    }
+
+    // SERVER-13789
+    TEST_F(QueryPlannerTest, ElemMatchIndexedNestedNE) {
+        addIndex(BSON("bar.baz" << 1));
+        runQuery(fromjson("{foo: 1, $and: [{bar: {$elemMatch: {baz: {$ne: 2}}}}]}"));
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{cscan: {dir: 1}}");
+        assertSolutionExists("{fetch: {filter: {$and: [{foo:1},"
+                                               "{bar:{$elemMatch:{baz:{$ne:2}}}}]}, "
+                                "node: {ixscan: {pattern: {'bar.baz': 1}, "
+                                                "bounds: {'bar.baz': [['MinKey',2,true,false], "
+                                                "[2,'MaxKey',false,true]]}}}}}");
+    }
+
+    // SERVER-13789: Make sure we properly handle an $or below $elemMatch that is not
+    // tagged by the enumerator to use an index.
+    TEST_F(QueryPlannerTest, ElemMatchNestedOrNotIndexed) {
+        addIndex(BSON("a.b" << 1));
+        runQuery(fromjson("{c: 1, a: {$elemMatch: {b: 3, $or: [{c: 4}, {c: 5}]}}}"));
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{cscan: {dir: 1}}");
+        assertSolutionExists("{fetch: {node: {ixscan: {pattern: {'a.b': 1}, bounds: "
+                                "{'a.b': [[3,3,true,true]]}}}}}");
+    }
+
     //
     // Geo
     // http://docs.mongodb.org/manual/reference/operator/query-geospatial/#geospatial-query-compatibility-chart