SERVER-15886 SubplanStage tries to plan each branch from the cache

Fixes a bug in which the SubplanStage could put a bad plan into the plan cache.

3 files changed, 270 insertions, 185 deletions

diff --git a/src/mongo/db/exec/subplan.cpp b/src/mongo/db/exec/subplan.cpp
index a83b346cc03..e436fe8e01a 100644
--- a/src/mongo/db/exec/subplan.cpp
+++ b/src/mongo/db/exec/subplan.cpp
@@ -30,11 +30,9 @@
 
 #include "mongo/db/exec/subplan.h"
 
-#include "mongo/base/owned_pointer_vector.h"
 #include "mongo/client/dbclientinterface.h"
 #include "mongo/db/exec/multi_plan.h"
 #include "mongo/db/exec/scoped_timer.h"
-#include "mongo/db/query/canonical_query.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/query/plan_executor.h"
 #include "mongo/db/query/planner_analysis.h"
@@ -58,25 +56,9 @@ namespace mongo {
           _ws(ws),
           _plannerParams(params),
           _query(cq),
-          _killed(false),
           _child(NULL),
           _commonStats(kStageType) { }
 
-    SubplanStage::~SubplanStage() {
-        while (!_solutions.empty()) {
-            vector<QuerySolution*> solns = _solutions.front();
-            for (size_t i = 0; i < solns.size(); i++) {
-                delete solns[i];
-            }
-            _solutions.pop();
-        }
-
-        while (!_cqs.empty()) {
-            delete _cqs.front();
-            _cqs.pop();
-        }
-    }
-
     // static
     bool SubplanStage::canUseSubplanning(const CanonicalQuery& query) {
         const LiteParsedQuery& lpq = query.getParsed();
@@ -87,13 +69,6 @@ namespace mongo {
             return false;
         }
 
-        // Collection scan
-        // No sort order requested
-        if (lpq.getSort().isEmpty() &&
-            expr->matchType() == MatchExpression::AND && expr->numChildren() == 0) {
-            return false;
-        }
-
         // Hint provided
         if (!lpq.getHint().isEmpty()) {
             return false;
@@ -129,7 +104,7 @@ namespace mongo {
         // work that happens here, so this is needed for the time accounting to make sense.
         ScopedTimer timer(&_commonStats.executionTimeMillis);
 
-        MatchExpression* theOr = _query->root();
+        MatchExpression* orExpr = _query->root();
 
         for (size_t i = 0; i < _plannerParams.indices.size(); ++i) {
             const IndexEntry& ie = _plannerParams.indices[i];
@@ -139,130 +114,178 @@ namespace mongo {
 
         const WhereCallbackReal whereCallback(_txn, _collection->ns().db());
 
-        for (size_t i = 0; i < theOr->numChildren(); ++i) {
+        for (size_t i = 0; i < orExpr->numChildren(); ++i) {
+            // We need a place to shove the results from planning this branch.
+            _branchResults.push_back(new BranchPlanningResult());
+            BranchPlanningResult* branchResult = _branchResults.back();
+
+            MatchExpression* orChild = orExpr->getChild(i);
+
             // Turn the i-th child into its own query.
-            MatchExpression* orChild = theOr->getChild(i);
-            CanonicalQuery* orChildCQ;
-            Status childCQStatus = CanonicalQuery::canonicalize(*_query,
-                                                                orChild,
-                                                                &orChildCQ,
-                                                                whereCallback);
-            if (!childCQStatus.isOK()) {
-                mongoutils::str::stream ss;
-                ss << "Can't canonicalize subchild " << orChild->toString()
-                   << " " << childCQStatus.reason();
-                return Status(ErrorCodes::BadValue, ss);
+            {
+                CanonicalQuery* orChildCQ;
+                Status childCQStatus = CanonicalQuery::canonicalize(*_query,
+                                                                    orChild,
+                                                                    &orChildCQ,
+                                                                    whereCallback);
+                if (!childCQStatus.isOK()) {
+                    mongoutils::str::stream ss;
+                    ss << "Can't canonicalize subchild " << orChild->toString()
+                       << " " << childCQStatus.reason();
+                    return Status(ErrorCodes::BadValue, ss);
+                }
+
+                branchResult->canonicalQuery.reset(orChildCQ);
             }
 
-            // Make sure it gets cleaned up.
-            auto_ptr<CanonicalQuery> safeOrChildCQ(orChildCQ);
+            // Plan the i-th child. We might be able to find a plan for the i-th child in the plan
+            // cache. If there's no cached plan, then we generate and rank plans using the MPS.
+            CachedSolution* rawCS;
+            if (PlanCache::shouldCacheQuery(*branchResult->canonicalQuery.get()) &&
+                _collection->infoCache()->getPlanCache()->get(*branchResult->canonicalQuery.get(),
+                                                              &rawCS).isOK()) {
+                // We have a CachedSolution. Store it for later.
+                QLOG() << "Subplanner: cached plan found for child " << i << " of "
+                       << orExpr->numChildren();
+
+                branchResult->cachedSolution.reset(rawCS);
+            }
+            else {
+                // No CachedSolution found. We'll have to plan from scratch.
+                QLOG() << "Subplanner: planning child " << i << " of " << orExpr->numChildren();
+
+                // We don't set NO_TABLE_SCAN because peeking at the cache data will keep us from
+                // considering any plan that's a collscan.
+                Status status = QueryPlanner::plan(*branchResult->canonicalQuery.get(),
+                                                   _plannerParams,
+                                                   &branchResult->solutions.mutableVector());
 
-            // Plan the i-th child.
-            vector<QuerySolution*> solutions;
+                if (!status.isOK()) {
+                    mongoutils::str::stream ss;
+                    ss << "Can't plan for subchild "
+                       << branchResult->canonicalQuery->toString()
+                       << " " << status.reason();
+                    return Status(ErrorCodes::BadValue, ss);
+                }
+                QLOG() << "Subplanner: got " << branchResult->solutions.size() << " solutions";
+
+                if (0 == branchResult->solutions.size()) {
+                    // If one child doesn't have an indexed solution, bail out.
+                    mongoutils::str::stream ss;
+                    ss << "No solutions for subchild " << branchResult->canonicalQuery->toString();
+                    return Status(ErrorCodes::BadValue, ss);
+                }
+            }
+        }
+
+        return Status::OK();
+    }
 
-            // We don't set NO_TABLE_SCAN because peeking at the cache data will keep us from 
-            // considering any plan that's a collscan.
-            QLOG() << "Subplanner: planning child " << i << " of " << theOr->numChildren();
-            Status status = QueryPlanner::plan(*safeOrChildCQ, _plannerParams, &solutions);
+    namespace {
+
+        /**
+         * On success, applies the index tags from 'branchCacheData' (which represent the winning
+         * plan for 'orChild') to 'compositeCacheData'.
+         */
+        Status tagOrChildAccordingToCache(PlanCacheIndexTree* compositeCacheData,
+                                          SolutionCacheData* branchCacheData,
+                                          MatchExpression* orChild,
+                                          const std::map<BSONObj, size_t>& indexMap) {
+            invariant(compositeCacheData);
+
+            // We want a well-formed *indexed* solution.
+            if (NULL == branchCacheData) {
+                // For example, we don't cache things for 2d indices.
+                mongoutils::str::stream ss;
+                ss << "No cache data for subchild " << orChild->toString();
+                return Status(ErrorCodes::BadValue, ss);
+            }
 
-            if (!status.isOK()) {
+            if (SolutionCacheData::USE_INDEX_TAGS_SOLN != branchCacheData->solnType) {
                 mongoutils::str::stream ss;
-                ss << "Can't plan for subchild " << orChildCQ->toString()
-                   << " " << status.reason();
+                ss << "No indexed cache data for subchild "
+                   << orChild->toString();
                 return Status(ErrorCodes::BadValue, ss);
             }
-            QLOG() << "Subplanner: got " << solutions.size() << " solutions";
 
-            if (0 == solutions.size()) {
-                // If one child doesn't have an indexed solution, bail out.
+            // Add the index assignments to our original query.
+            Status tagStatus = QueryPlanner::tagAccordingToCache(orChild,
+                                                                 branchCacheData->tree.get(),
+                                                                 indexMap);
+
+            if (!tagStatus.isOK()) {
                 mongoutils::str::stream ss;
-                ss << "No solutions for subchild " << orChildCQ->toString();
+                ss << "Failed to extract indices from subchild "
+                   << orChild->toString();
                 return Status(ErrorCodes::BadValue, ss);
             }
 
-            // Hang onto the canonicalized subqueries and the corresponding query solutions
-            // so that they can be used in subplan running later on.
-            _cqs.push(safeOrChildCQ.release());
-            _solutions.push(solutions);
+            // Add the child's cache data to the cache data we're creating for the main query.
+            compositeCacheData->children.push_back(branchCacheData->tree->clone());
+
+            return Status::OK();
         }
 
-        return Status::OK();
-    }
+    } // namespace
 
     Status SubplanStage::choosePlanForSubqueries(PlanYieldPolicy* yieldPolicy) {
         // This is what we annotate with the index selections and then turn into a solution.
-        auto_ptr<OrMatchExpression> theOr(
+        auto_ptr<OrMatchExpression> orExpr(
             static_cast<OrMatchExpression*>(_query->root()->shallowClone()));
 
         // This is the skeleton of index selections that is inserted into the cache.
         auto_ptr<PlanCacheIndexTree> cacheData(new PlanCacheIndexTree());
 
-        for (size_t i = 0; i < theOr->numChildren(); ++i) {
-            MatchExpression* orChild = theOr->getChild(i);
-
-            auto_ptr<CanonicalQuery> orChildCQ(_cqs.front());
-            _cqs.pop();
-
-            // 'solutions' is owned by the SubplanStage instance until
-            // it is popped from the queue.
-            vector<QuerySolution*> solutions = _solutions.front();
-            _solutions.pop();
-
-            // We already checked for zero solutions in planSubqueries(...).
-            invariant(!solutions.empty());
-
-            if (1 == solutions.size()) {
-                // There is only one solution. Transfer ownership to an auto_ptr.
-                auto_ptr<QuerySolution> autoSoln(solutions[0]);
-
-                // We want a well-formed *indexed* solution.
-                if (NULL == autoSoln->cacheData.get()) {
-                    // For example, we don't cache things for 2d indices.
-                    mongoutils::str::stream ss;
-                    ss << "No cache data for subchild " << orChild->toString();
-                    return Status(ErrorCodes::BadValue, ss);
-                }
-
-                if (SolutionCacheData::USE_INDEX_TAGS_SOLN != autoSoln->cacheData->solnType) {
-                    mongoutils::str::stream ss;
-                    ss << "No indexed cache data for subchild "
-                       << orChild->toString();
-                    return Status(ErrorCodes::BadValue, ss);
+        for (size_t i = 0; i < orExpr->numChildren(); ++i) {
+            MatchExpression* orChild = orExpr->getChild(i);
+            BranchPlanningResult* branchResult = _branchResults[i];
+
+            if (branchResult->cachedSolution.get()) {
+                // We can get the index tags we need out of the cache.
+                Status tagStatus = tagOrChildAccordingToCache(
+                    cacheData.get(),
+                    branchResult->cachedSolution->plannerData[0],
+                    orChild,
+                    _indexMap);
+                if (!tagStatus.isOK()) {
+                    return tagStatus;
                 }
-
-                // Add the index assignments to our original query.
-                Status tagStatus = QueryPlanner::tagAccordingToCache(
-                    orChild, autoSoln->cacheData->tree.get(), _indexMap);
-
+            }
+            else if (1 == branchResult->solutions.size()) {
+                QuerySolution* soln = branchResult->solutions.front();
+                Status tagStatus = tagOrChildAccordingToCache(cacheData.get(),
+                                                              soln->cacheData.get(),
+                                                              orChild,
+                                                              _indexMap);
                 if (!tagStatus.isOK()) {
-                    mongoutils::str::stream ss;
-                    ss << "Failed to extract indices from subchild "
-                       << orChild->toString();
-                    return Status(ErrorCodes::BadValue, ss);
+                    return tagStatus;
                 }
-
-                // Add the child's cache data to the cache data we're creating for the main query.
-                cacheData->children.push_back(autoSoln->cacheData->tree->clone());
             }
             else {
-                // N solutions, rank them.  Takes ownership of orChildCQ.
+                // N solutions, rank them.
+
+                // We already checked for zero solutions in planSubqueries(...).
+                invariant(!branchResult->solutions.empty());
 
                 _ws->clear();
 
-                _child.reset(new MultiPlanStage(_txn, _collection, orChildCQ.get()));
+                _child.reset(new MultiPlanStage(_txn, _collection,
+                                                branchResult->canonicalQuery.get()));
                 MultiPlanStage* multiPlanStage = static_cast<MultiPlanStage*>(_child.get());
 
-                // Dump all the solutions into the MPR.
-                for (size_t ix = 0; ix < solutions.size(); ++ix) {
+                // Dump all the solutions into the MPS.
+                for (size_t ix = 0; ix < branchResult->solutions.size(); ++ix) {
                     PlanStage* nextPlanRoot;
-                    verify(StageBuilder::build(_txn,
-                                               _collection,
-                                               *solutions[ix],
-                                               _ws,
-                                               &nextPlanRoot));
-
-                    // Owns first two arguments
-                    multiPlanStage->addPlan(solutions[ix], nextPlanRoot, _ws);
+                    invariant(StageBuilder::build(_txn,
+                                                  _collection,
+                                                  *branchResult->solutions[ix],
+                                                  _ws,
+                                                  &nextPlanRoot));
+
+                    // Takes ownership of solution with index 'ix' and 'nextPlanRoot'.
+                    multiPlanStage->addPlan(branchResult->solutions.releaseAt(ix),
+                                            nextPlanRoot,
+                                            _ws);
                 }
 
                 Status planSelectStat = multiPlanStage->pickBestPlan(yieldPolicy);
@@ -273,7 +296,7 @@ namespace mongo {
                 if (!multiPlanStage->bestPlanChosen()) {
                     mongoutils::str::stream ss;
                     ss << "Failed to pick best plan for subchild "
-                       << orChildCQ->toString();
+                       << branchResult->canonicalQuery->toString();
                     return Status(ErrorCodes::BadValue, ss);
                 }
 
@@ -310,11 +333,11 @@ namespace mongo {
         }
 
         // Must do this before using the planner functionality.
-        sortUsingTags(theOr.get());
+        sortUsingTags(orExpr.get());
 
-        // Use the cached index assignments to build solnRoot.  Takes ownership of 'theOr'
+        // Use the cached index assignments to build solnRoot.  Takes ownership of 'orExpr'.
         QuerySolutionNode* solnRoot = QueryPlannerAccess::buildIndexedDataAccess(
-            *_query, theOr.release(), false, _plannerParams.indices);
+            *_query, orExpr.release(), false, _plannerParams.indices);
 
         if (NULL == solnRoot) {
             mongoutils::str::stream ss;
@@ -325,40 +348,24 @@ namespace mongo {
         QLOG() << "Subplanner: fully tagged tree is " << solnRoot->toString();
 
         // Takes ownership of 'solnRoot'
-        QuerySolution* soln = QueryPlannerAnalysis::analyzeDataAccess(*_query,
-                                                                      _plannerParams,
-                                                                      solnRoot);
+        _compositeSolution.reset(QueryPlannerAnalysis::analyzeDataAccess(*_query,
+                                                                         _plannerParams,
+                                                                         solnRoot));
 
-        if (NULL == soln) {
+        if (NULL == _compositeSolution.get()) {
             mongoutils::str::stream ss;
             ss << "Failed to analyze subplanned query";
             return Status(ErrorCodes::BadValue, ss);
         }
 
-        // We want our franken-solution to be cached.
-        SolutionCacheData* scd = new SolutionCacheData();
-        scd->tree.reset(cacheData.release());
-        soln->cacheData.reset(scd);
-
-        QLOG() << "Subplanner: Composite solution is " << soln->toString() << endl;
+        QLOG() << "Subplanner: Composite solution is " << _compositeSolution->toString() << endl;
 
-        // We use one of these even if there is one plan.  We do this so that the entry is cached
-        // with stats obtained in the same fashion as a competitive ranking would have obtained
-        // them.
+        // Use the index tags from planning each branch to construct the composite solution,
+        // and set that solution as our child stage.
         _ws->clear();
-        _child.reset(new MultiPlanStage(_txn, _collection, _query));
-        MultiPlanStage* multiPlanStage = static_cast<MultiPlanStage*>(_child.get());
         PlanStage* root;
-        verify(StageBuilder::build(_txn, _collection, *soln, _ws, &root));
-        multiPlanStage->addPlan(soln, root, _ws); // Takes ownership first two arguments.
-
-        multiPlanStage->pickBestPlan(yieldPolicy);
-        if (! multiPlanStage->bestPlanChosen()) {
-            mongoutils::str::stream ss;
-            ss << "Failed to pick best plan for subchild "
-               << _query->toString();
-            return Status(ErrorCodes::BadValue, ss);
-        }
+        invariant(StageBuilder::build(_txn, _collection, *_compositeSolution.get(), _ws, &root));
+        _child.reset(root);
 
         return Status::OK();
     }
@@ -395,7 +402,7 @@ namespace mongo {
             _child.reset(root);
 
             // This SubplanStage takes ownership of the query solution.
-            _solutions.push(solutions.release());
+            _compositeSolution.reset(solutions.popAndReleaseBack());
 
             return Status::OK();
         }
@@ -452,10 +459,6 @@ namespace mongo {
     }
 
     bool SubplanStage::isEOF() {
-        if (_killed) {
-            return true;
-        }
-
         // If we're running we best have a runner.
         invariant(_child.get());
         return _child->isEOF();
@@ -467,10 +470,6 @@ namespace mongo {
         // Adds the amount of time taken by work() to executionTimeMillis.
         ScopedTimer timer(&_commonStats.executionTimeMillis);
 
-        if (_killed) {
-            return PlanStage::DEAD;
-        }
-
         if (isEOF()) { return PlanStage::IS_EOF; }
 
         invariant(_child.get());
@@ -492,11 +491,8 @@ namespace mongo {
     void SubplanStage::saveState() {
         _txn = NULL;
         ++_commonStats.yields;
-        if (_killed) {
-            return;
-        }
 
-        // We're ranking a sub-plan via an MPR or we're streaming results from this stage.  Either
+        // We're ranking a sub-plan via an MPS or we're streaming results from this stage.  Either
         // way, pass on the request.
         if (NULL != _child.get()) {
             _child->saveState();
@@ -507,11 +503,8 @@ namespace mongo {
         invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
-        if (_killed) {
-            return;
-        }
 
-        // We're ranking a sub-plan via an MPR or we're streaming results from this stage.  Either
+        // We're ranking a sub-plan via an MPS or we're streaming results from this stage.  Either
         // way, pass on the request.
         if (NULL != _child.get()) {
             _child->restoreState(opCtx);
@@ -520,9 +513,6 @@ namespace mongo {
 
     void SubplanStage::invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type) {
         ++_commonStats.invalidates;
-        if (_killed) {
-            return;
-        }
 
         if (NULL != _child.get()) {
             _child->invalidate(txn, dl, type);
@@ -544,6 +534,10 @@ namespace mongo {
         return ret.release();
     }
 
+    bool SubplanStage::branchPlannedFromCache(size_t i) const {
+        return NULL != _branchResults[i]->cachedSolution.get();
+    }
+
     const CommonStats* SubplanStage::getCommonStats() {
         return &_commonStats;
     }
diff --git a/src/mongo/db/exec/subplan.h b/src/mongo/db/exec/subplan.h
index c2ae5b78ed1..ec275c66a92 100644
--- a/src/mongo/db/exec/subplan.h
+++ b/src/mongo/db/exec/subplan.h
@@ -30,10 +30,12 @@
 
 #include <boost/scoped_ptr.hpp>
 #include <string>
-#include <queue>
 
+#include "mongo/base/owned_pointer_vector.h"
 #include "mongo/base/status.h"
 #include "mongo/db/exec/plan_stage.h"
+#include "mongo/db/query/canonical_query.h"
+#include "mongo/db/query/plan_cache.h"
 #include "mongo/db/query/plan_yield_policy.h"
 #include "mongo/db/query/query_planner_params.h"
 #include "mongo/db/query/query_solution.h"
@@ -49,6 +51,18 @@ namespace mongo {
      * each clause.
      *
      * Uses the MultiPlanStage in order to rank plans for the individual clauses.
+     *
+     * Notes on caching strategy:
+     *
+     *   --Interaction with the plan cache is done on a per-clause basis. For a given clause C,
+     *   if there is a plan in the cache for shape C, then C is planned using the index tags
+     *   obtained from the plan cache entry. If no cached plan is found for C, then a MultiPlanStage
+     *   is used to determine the best plan for the clause; unless there is a tie between multiple
+     *   candidate plans, the winner is inserted into the plan cache and used to plan subsequent
+     *   executions of C. These subsequent executions of shape C could be either as a clause in
+     *   another rooted $or query, or shape C as its own query.
+     *
+     *   --Plans for entire rooted $or queries are neither written to nor read from the plan cache.
      */
     class SubplanStage : public PlanStage {
     public:
@@ -58,8 +72,6 @@ namespace mongo {
                      const QueryPlannerParams& params,
                      CanonicalQuery* cq);
 
-        virtual ~SubplanStage();
-
         static bool canUseSubplanning(const CanonicalQuery& query);
 
         virtual bool isEOF();
@@ -96,8 +108,40 @@ namespace mongo {
          */
         Status pickBestPlan(PlanYieldPolicy* yieldPolicy);
 
+        //
+        // For testing.
+        //
+
+        /**
+         * Returns true if the i-th branch was planned by retrieving a cached solution,
+         * otherwise returns false.
+         */
+        bool branchPlannedFromCache(size_t i) const;
+
     private:
         /**
+         * A class used internally in order to keep track of the results of planning
+         * a particular $or branch.
+         */
+        struct BranchPlanningResult {
+            MONGO_DISALLOW_COPYING(BranchPlanningResult);
+        public:
+            BranchPlanningResult() { }
+
+            // A parsed version of one branch of the $or.
+            boost::scoped_ptr<CanonicalQuery> canonicalQuery;
+
+            // If there is cache data available, then we store it here rather than generating
+            // a set of alternate plans for the branch. The index tags from the cache data
+            // can be applied directly to the parent $or MatchExpression when generating the
+            // composite solution.
+            boost::scoped_ptr<CachedSolution> cachedSolution;
+
+            // Query solutions resulting from planning the $or branch.
+            OwnedPointerVector<QuerySolution> solutions;
+        };
+
+        /**
          * Plan each branch of the $or independently, and store the resulting
          * lists of query solutions in '_solutions'.
          *
@@ -133,24 +177,17 @@ namespace mongo {
         // Not owned here.
         CanonicalQuery* _query;
 
-        bool _killed;
+        // If we successfully create a "composite solution" by planning each $or branch
+        // independently, that solution is owned here.
+        boost::scoped_ptr<QuerySolution> _compositeSolution;
 
         boost::scoped_ptr<PlanStage> _child;
 
-        // We do the subquery planning up front, and keep the resulting query solutions here. Lists
-        // of query solutions are dequeued and ownership is transferred to the underlying
-        // MultiPlanStages one at a time.
-        //
-        // If we fall back on regular planning and find that there is only a single query solution,
-        // then the SubplanStage retains ownership of that solution here.
-        std::queue< std::vector<QuerySolution*> > _solutions;
-
-        // Holds the canonicalized subqueries. Ownership is transferred
-        // to the underlying runners one at a time.
-        std::queue<CanonicalQuery*> _cqs;
+        // Holds a list of the results from planning each branch.
+        OwnedPointerVector<BranchPlanningResult> _branchResults;
 
         // We need this to extract cache-friendly index data from the index assignments.
-        map<BSONObj, size_t> _indexMap;
+        std::map<BSONObj, size_t> _indexMap;
 
         CommonStats _commonStats;
     };
diff --git a/src/mongo/dbtests/query_stage_subplan.cpp b/src/mongo/dbtests/query_stage_subplan.cpp
index 18cc6e9b2f9..309e7516952 100644
--- a/src/mongo/dbtests/query_stage_subplan.cpp
+++ b/src/mongo/dbtests/query_stage_subplan.cpp
@@ -82,19 +82,19 @@ namespace QueryStageSubplan {
             BSONObj query = fromjson("{$or: [{a: {$geoWithin: {$centerSphere: [[0,0],10]}}},"
                                             "{a: {$geoWithin: {$centerSphere: [[1,1],10]}}}]}");
 
-            CanonicalQuery* cq;
-            ASSERT_OK(CanonicalQuery::canonicalize(ns(), query, &cq));
-            boost::scoped_ptr<CanonicalQuery> killCq(cq);
+            CanonicalQuery* rawCq;
+            ASSERT_OK(CanonicalQuery::canonicalize(ns(), query, &rawCq));
+            boost::scoped_ptr<CanonicalQuery> cq(rawCq);
 
             Collection* collection = ctx.getCollection();
 
             // Get planner params.
             QueryPlannerParams plannerParams;
-            fillOutPlannerParams(&_txn, collection, cq, &plannerParams);
+            fillOutPlannerParams(&_txn, collection, cq.get(), &plannerParams);
 
             WorkingSet ws;
             boost::scoped_ptr<SubplanStage> subplan(new SubplanStage(&_txn, collection, &ws,
-                                                                     plannerParams, cq));
+                                                                     plannerParams, cq.get()));
 
             // NULL means that 'subplan' will not yield during plan selection. Plan selection
             // should succeed due to falling back on regular planning.
@@ -102,12 +102,66 @@ namespace QueryStageSubplan {
         }
     };
 
+    /**
+     * Test the SubplanStage's ability to plan an individual branch using the plan cache.
+     */
+    class QueryStageSubplanPlanFromCache : public QueryStageSubplanBase {
+    public:
+        void run() {
+            Client::WriteContext ctx(&_txn, ns());
+
+            addIndex(BSON("a" << 1 << "b" << 1));
+            addIndex(BSON("a" << 1 << "c" << 1));
+
+            for (int i = 0; i < 10; i++) {
+                insert(BSON("a" << 1 << "b" << i << "c" << i));
+            }
+
+            // This query should result in a plan cache entry for the first branch. The second
+            // branch should tie, meaning that nothing is inserted into the plan cache.
+            BSONObj query = fromjson("{$or: [{a: 1, b: 3}, {a: 1}]}");
+
+            Collection* collection = ctx.getCollection();
+
+            CanonicalQuery* rawCq;
+            ASSERT_OK(CanonicalQuery::canonicalize(ns(), query, &rawCq));
+            boost::scoped_ptr<CanonicalQuery> cq(rawCq);
+
+            // Get planner params.
+            QueryPlannerParams plannerParams;
+            fillOutPlannerParams(&_txn, collection, cq.get(), &plannerParams);
+
+            WorkingSet ws;
+            boost::scoped_ptr<SubplanStage> subplan(new SubplanStage(&_txn, collection, &ws,
+                                                                     plannerParams, cq.get()));
+
+            // NULL means that 'subplan' should not yield during plan selection.
+            ASSERT_OK(subplan->pickBestPlan(NULL));
+
+            // Nothing is in the cache yet, so neither branch should have been planned from
+            // the plan cache.
+            ASSERT_FALSE(subplan->branchPlannedFromCache(0));
+            ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+
+            // If we repeat the same query, then the first branch should come from the cache,
+            // but the second is re-planned due to tying on the first run.
+            ws.clear();
+            subplan.reset(new SubplanStage(&_txn, collection, &ws, plannerParams, cq.get()));
+
+            ASSERT_OK(subplan->pickBestPlan(NULL));
+
+            ASSERT_TRUE(subplan->branchPlannedFromCache(0));
+            ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+        }
+    };
+
     class All : public Suite {
     public:
         All() : Suite("query_stage_subplan") {}
 
         void setupTests() {
             add<QueryStageSubplanGeo2dOr>();
+            add<QueryStageSubplanPlanFromCache>();
         }
     };