SERVER-19996 Keep Mutations no longer added on ntoreturn hack

(cherry picked from commit e07a93a7ea6756434e0d3df38cdd1e6f5d5fd010)

Conflicts:
	src/mongo/db/query/planner_analysis.cpp

2 files changed, 24 insertions, 4 deletions

diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 24a4e446b01..f21d8c79ca3 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -583,19 +583,24 @@ QuerySolution* QueryPlannerAnalysis::analyzeDataAccess(const CanonicalQuery& que
     //
     // 3. There is an index-provided sort.  Ditto above comment about merging.
     //
+    // 4. There is a SORT that is not at the root of solution tree. Ditto above comment about
+    // merging.
+    //
     // TODO: do we want some kind of pre-planning step where we look for certain nodes and cache
     // them?  We do lookups in the tree a few times.  This may not matter as most trees are
     // shallow in terms of query nodes.
-    bool cannotKeepFlagged = hasNode(solnRoot, STAGE_TEXT) ||
+    const bool hasNotRootSort = hasSortStage && STAGE_SORT != solnRoot->getType();
+
+    const bool cannotKeepFlagged = hasNode(solnRoot, STAGE_TEXT) ||
         hasNode(solnRoot, STAGE_GEO_NEAR_2D) || hasNode(solnRoot, STAGE_GEO_NEAR_2DSPHERE) ||
-        (!query.getParsed().getSort().isEmpty() && !hasSortStage);
+        (!query.getParsed().getSort().isEmpty() && !hasSortStage) || hasNotRootSort;
 
     // Only these stages can produce flagged results.  A stage has to hold state past one call
     // to work(...) in order to possibly flag a result.
-    bool couldProduceFlagged =
+    const bool couldProduceFlagged =
         hasAndHashStage || hasNode(solnRoot, STAGE_AND_SORTED) || hasNode(solnRoot, STAGE_FETCH);
 
-    bool shouldAddMutation = !cannotKeepFlagged && couldProduceFlagged;
+    const bool shouldAddMutation = !cannotKeepFlagged && couldProduceFlagged;
 
     if (shouldAddMutation && (params.options & QueryPlannerParams::KEEP_MUTATIONS)) {
         KeepMutationsNode* keep = new KeepMutationsNode();
diff --git a/src/mongo/db/query/query_planner_test.cpp b/src/mongo/db/query/query_planner_test.cpp
index d3378221acc..ee4b511f4e4 100644
--- a/src/mongo/db/query/query_planner_test.cpp
+++ b/src/mongo/db/query/query_planner_test.cpp
@@ -5423,6 +5423,21 @@ TEST_F(QueryPlannerTest, NoKeepWithIndexedSort) {
         "[{ixscan: {pattern: {a: 1, b: 1}}}, {ixscan: {pattern: {a: 1, b: 1}}}]}}}}");
 }
 
+// No KeepMutations when we have a sort that is not root, like the ntoreturn hack.
+TEST_F(QueryPlannerTest, NoKeepWithNToReturn) {
+    params.options = QueryPlannerParams::KEEP_MUTATIONS;
+    params.options |= QueryPlannerParams::SPLIT_LIMITED_SORT;
+    addIndex(BSON("a" << 1));
+    runQuerySortProjSkipLimit(fromjson("{a: 1}"), fromjson("{b: 1}"), BSONObj(), 0, 3);
+
+    assertSolutionExists(
+        "{or: {nodes: ["
+        "{sort: {pattern: {b: 1}, limit: 3, node: "
+        "{fetch: {node: {ixscan: {pattern: {a: 1}}}}}}}, "
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{fetch: {node: {ixscan: {pattern: {a: 1}}}}}}}]}}");
+}
+
 // Make sure a top-level $or hits the limiting number
 // of solutions that we are willing to consider.
 TEST_F(QueryPlannerTest, OrEnumerationLimit) {