SERVER-13754 explode OR solutions for sort

(cherry picked from commit 4fd34d5272e6b282f5b98b7b65c4ba2c092d93fb)

2 files changed, 134 insertions, 23 deletions

diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 888ff143cf2..0888cce7545 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -80,24 +80,39 @@ namespace mongo {
 
         /**
          * Should we try to expand the index scan(s) in 'solnRoot' to pull out an indexed sort?
+         *
+         * Returns the node which should be replaced by the merge sort of exploded scans
+         * in the out-parameter 'toReplace'.
          */
-        bool structureOKForExplode(QuerySolutionNode* solnRoot) {
+        bool structureOKForExplode(QuerySolutionNode* solnRoot, QuerySolutionNode** toReplace) {
             // For now we only explode if we *know* we will pull the sort out.  We can look at
             // more structure (or just explode and recalculate properties and see what happens)
             // but for now we just explode if it's a sure bet.
             //
-            // TODO: Can also try exploding if root is OR and children are ixscans, or root is
-            // AND_HASH (last child dictates order.), or other less obvious cases...
+            // TODO: Can also try exploding if root is AND_HASH (last child dictates order.),
+            // or other less obvious cases...
             if (STAGE_IXSCAN == solnRoot->getType()) {
+                *toReplace = solnRoot;
                 return true;
             }
 
             if (STAGE_FETCH == solnRoot->getType()) {
                  if (STAGE_IXSCAN == solnRoot->children[0]->getType()) {
+                     *toReplace = solnRoot->children[0];
                      return true;
                  }
             }
 
+            if (STAGE_OR == solnRoot->getType()) {
+                for (size_t i = 0; i < solnRoot->children.size(); ++i) {
+                    if (STAGE_IXSCAN != solnRoot->children[i]->getType()) {
+                        return false;
+                    }
+                }
+                *toReplace = solnRoot;
+                return true;
+            }
+
             return false;
         }
 
@@ -151,8 +166,9 @@ namespace mongo {
         }
 
         /**
-         * Take the provided index scan node 'isn' and return a logically equivalent node that
-         * provides the same data but provides the sort order 'sort'.
+         * Take the provided index scan node 'isn'. Returns a list of index scans which are
+         * logically equivalent to 'isn' if joined by a MergeSort through the out-parameter
+         * 'explosionResult'. These index scan instances are owned by the caller.
          *
          * fieldsToExplode is a count of how many fields in the scan's bounds are the union of point
          * intervals.  This is computed beforehand and provided as a small optimization.
@@ -164,22 +180,19 @@ namespace mongo {
          * 'sort' will be {b: 1}
          * 'fieldsToExplode' will be 1 (as only one field isUnionOfPoints).
          *
-         * The solution returned will be a mergesort of the two scans:
+         * On return, 'explosionResult' will contain the following two scans:
          * a:[[1,1]], b:[MinKey, MaxKey]
          * a:[[2,2]], b:[MinKey, MaxKey]
          */
-        QuerySolutionNode* explodeScan(IndexScanNode* isn,
-                                       const BSONObj& sort,
-                                       size_t fieldsToExplode) {
+        void explodeScan(IndexScanNode* isn,
+                         const BSONObj& sort,
+                         size_t fieldsToExplode,
+                         vector<QuerySolutionNode*>* explosionResult) {
 
             // Turn the compact bounds in 'isn' into a bunch of points...
             vector<PointPrefix> prefixForScans;
             makeCartesianProduct(isn->bounds, fieldsToExplode, &prefixForScans);
 
-            // And merge-sort the scans over those points.
-            auto_ptr<MergeSortNode> merge(new MergeSortNode());
-            merge->sort = sort;
-
             for (size_t i = 0; i < prefixForScans.size(); ++i) {
                 const PointPrefix& prefix = prefixForScans[i];
                 verify(prefix.size() == fieldsToExplode);
@@ -201,11 +214,8 @@ namespace mongo {
                 for (size_t j = fieldsToExplode; j < isn->bounds.fields.size(); ++j) {
                     child->bounds.fields[j] = isn->bounds.fields[j];
                 }
-                merge->children.push_back(child);
+                explosionResult->push_back(child);
             }
-
-            merge->computeProperties();
-            return merge.release();
         }
 
         /**
@@ -245,7 +255,8 @@ namespace mongo {
                                               QuerySolutionNode** solnRoot) {
         vector<QuerySolutionNode*> leafNodes;
 
-        if (!structureOKForExplode(*solnRoot)) {
+        QuerySolutionNode* toReplace;
+        if (!structureOKForExplode(*solnRoot, &toReplace)) {
             return false;
         }
 
@@ -336,15 +347,20 @@ namespace mongo {
 
         // If we're here, we can (probably?  depends on how restrictive the structure check is)
         // get our sort order via ixscan blow-up.
+        MergeSortNode* merge = new MergeSortNode();
+        merge->sort = desiredSort;
         for (size_t i = 0; i < leafNodes.size(); ++i) {
             IndexScanNode* isn = static_cast<IndexScanNode*>(leafNodes[i]);
-            QuerySolutionNode* newNode = explodeScan(isn, desiredSort, fieldsToExplode[i]);
-            // Replace 'isn' with 'newNode'
-            replaceNodeInTree(solnRoot, isn, newNode);
-            // And get rid of the old data access node.
-            delete isn;
+            explodeScan(isn, desiredSort, fieldsToExplode[i], &merge->children);
         }
 
+        merge->computeProperties();
+
+        // Replace 'toReplace' with the new merge sort node.
+        replaceNodeInTree(solnRoot, toReplace, merge);
+        // And get rid of the node that got replaced.
+        delete toReplace;
+
         return true;
     }
 
diff --git a/src/mongo/db/query/query_planner_test.cpp b/src/mongo/db/query/query_planner_test.cpp
index 4ca7e6c4b86..5da03f05540 100644
--- a/src/mongo/db/query/query_planner_test.cpp
+++ b/src/mongo/db/query/query_planner_test.cpp
@@ -2093,6 +2093,101 @@ namespace {
                                 "{fetch: {node: {ixscan: {pattern: {a:1,b:1,c:1}}}}}}}");
     }
 
+    // SERVER-13754: exploding an $or
+    TEST_F(QueryPlannerTest, ExplodeOrForSort) {
+        addIndex(BSON("a" << 1 << "c" << 1));
+        addIndex(BSON("b" << 1 << "c" << 1));
+
+        runQuerySortProj(fromjson("{$or: [{a: 1}, {a: 2}, {b: 2}]}"),
+                         BSON("c" << 1), BSONObj());
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{sort: {pattern: {c: 1}, limit: 0, node: {cscan: {dir: 1}}}}");
+        assertSolutionExists("{fetch: {node: {mergeSort: {nodes: "
+                                "[{ixscan: {bounds: {a: [[1,1,true,true]], "
+                                                    "c: [['MinKey','MaxKey',true,true]]},"
+                                           "pattern: {a:1, c:1}}},"
+                                 "{ixscan: {bounds: {a: [[2,2,true,true]], "
+                                                    "c: [['MinKey','MaxKey',true,true]]},"
+                                           "pattern: {a:1, c:1}}},"
+                                 "{ixscan: {bounds: {b: [[2,2,true,true]], "
+                                                    "c: [['MinKey','MaxKey',true,true]]},"
+                                           "pattern: {b:1, c:1}}}]}}}}");
+    }
+
+    // SERVER-13754: exploding an $or
+    TEST_F(QueryPlannerTest, ExplodeOrForSort2) {
+        addIndex(BSON("a" << 1 << "b" << 1 << "c" << 1));
+        addIndex(BSON("d" << 1 << "c" << 1));
+
+        runQuerySortProj(fromjson("{$or: [{a: 1, b: {$in: [1, 2]}}, {d: 3}]}"),
+                         BSON("c" << 1), BSONObj());
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{sort: {pattern: {c: 1}, limit: 0, node: {cscan: {dir: 1}}}}");
+        assertSolutionExists("{fetch: {node: {mergeSort: {nodes: "
+                                "[{ixscan: {bounds: {a: [[1,1,true,true]], b: [[1,1,true,true]],"
+                                                    "c: [['MinKey','MaxKey',true,true]]},"
+                                           "pattern: {a:1, b:1, c:1}}},"
+                                 "{ixscan: {bounds: {a: [[1,1,true,true]], b: [[2,2,true,true]],"
+                                                    "c: [['MinKey','MaxKey',true,true]]},"
+                                           "pattern: {a:1, b:1, c:1}}},"
+                                 "{ixscan: {bounds: {d: [[3,3,true,true]], "
+                                                    "c: [['MinKey','MaxKey',true,true]]},"
+                                           "pattern: {d:1, c:1}}}]}}}}");
+    }
+
+    // SERVER-13754: an $or that can't be exploded, because one clause of the
+    // $or does provide the sort, even after explosion.
+    TEST_F(QueryPlannerTest, CantExplodeOrForSort) {
+        addIndex(BSON("a" << 1 << "b" << 1 << "c" << 1));
+        addIndex(BSON("d" << 1 << "c" << 1));
+
+        runQuerySortProj(fromjson("{$or: [{a: {$in: [1, 2]}}, {d: 3}]}"),
+                         BSON("c" << 1), BSONObj());
+
+        assertNumSolutions(2U);
+        assertSolutionExists("{sort: {pattern: {c: 1}, limit: 0, node: {cscan: {dir: 1}}}}");
+        assertSolutionExists("{sort: {pattern: {c: 1}, limit: 0, node: "
+                                "{fetch: {filter: null, node: {or: {nodes: ["
+                                    "{ixscan: {pattern: {a: 1, b: 1, c: 1}}},"
+                                    "{ixscan: {pattern: {d: 1, c: 1}}}]}}}}}}");
+    }
+
+    // SERVER-13754: too many scans in an $or explosion.
+    TEST_F(QueryPlannerTest, TooManyToExplodeOr) {
+        addIndex(BSON("a" << 1 << "e" << 1));
+        addIndex(BSON("b" << 1 << "e" << 1));
+        addIndex(BSON("c" << 1 << "e" << 1));
+        addIndex(BSON("d" << 1 << "e" << 1));
+        runQuerySortProj(fromjson("{$or: [{a: {$in: [1,2,3,4,5,6]},"
+                                          "b: {$in: [1,2,3,4,5,6]}},"
+                                         "{c: {$in: [1,2,3,4,5,6]},"
+                                          "d: {$in: [1,2,3,4,5,6]}}]}"),
+                         BSON("e" << 1), BSONObj());
+
+        // We cap the # of ixscans we're willing to create, so we don't get explosion. Instead
+        // we get 5 different solutions which all use a blocking sort.
+        assertNumSolutions(5U);
+        assertSolutionExists("{sort: {pattern: {e: 1}, limit: 0, node: {cscan: {dir: 1}}}}");
+        assertSolutionExists("{sort: {pattern: {e: 1}, limit: 0, node: "
+                                "{or: {nodes: ["
+                                    "{fetch: {node: {ixscan: {pattern: {a: 1, e: 1}}}}},"
+                                    "{fetch: {node: {ixscan: {pattern: {c: 1, e: 1}}}}}]}}}}");
+        assertSolutionExists("{sort: {pattern: {e: 1}, limit: 0, node: "
+                                "{or: {nodes: ["
+                                    "{fetch: {node: {ixscan: {pattern: {b: 1, e: 1}}}}},"
+                                    "{fetch: {node: {ixscan: {pattern: {c: 1, e: 1}}}}}]}}}}");
+        assertSolutionExists("{sort: {pattern: {e: 1}, limit: 0, node: "
+                                "{or: {nodes: ["
+                                    "{fetch: {node: {ixscan: {pattern: {a: 1, e: 1}}}}},"
+                                    "{fetch: {node: {ixscan: {pattern: {d: 1, e: 1}}}}}]}}}}");
+        assertSolutionExists("{sort: {pattern: {e: 1}, limit: 0, node: "
+                                "{or: {nodes: ["
+                                    "{fetch: {node: {ixscan: {pattern: {b: 1, e: 1}}}}},"
+                                    "{fetch: {node: {ixscan: {pattern: {d: 1, e: 1}}}}}]}}}}");
+    }
+
     TEST_F(QueryPlannerTest, InWithSortAndLimitTrailingField) {
         addIndex(BSON("a" << 1 << "b" << -1 << "c" << 1));
         runQuerySortProjSkipLimit(fromjson("{a: {$in: [1, 2]}, b: {$gte: 0}}"),