SERVER-31858 Fix explodeForSort() path to handle multikeyness correctly.

5 files changed, 89 insertions, 2 deletions

diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 5b72f1c98fa..cd11da91262 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -373,6 +373,12 @@ bool QueryPlannerAnalysis::explodeForSort(const CanonicalQuery& query,
             return false;
         }
 
+        if (isn->index.multikey && isn->index.multikeyPaths.empty()) {
+            // The index is multikey but has no path-level multikeyness metadata. In this case, the
+            // index can never provide a sort.
+            return false;
+        }
+
         // How many scans will we create if we blow up this ixscan?
         size_t numScans = 1;
 
@@ -405,7 +411,13 @@ bool QueryPlannerAnalysis::explodeForSort(const CanonicalQuery& query,
         // the bounds.
         BSONObjBuilder resultingSortBob;
         while (kpIt.more()) {
-            resultingSortBob.append(kpIt.next());
+            auto elem = kpIt.next();
+            if (isn->multikeyFields.find(elem.fieldNameStringData()) != isn->multikeyFields.end()) {
+                // One of the indexed fields providing the sort is multikey. It is not correct for a
+                // field with multikey components to provide a sort, so bail out.
+                return false;
+            }
+            resultingSortBob.append(elem);
         }
 
         // See if it's the order we're looking for.
diff --git a/src/mongo/db/query/query_planner_array_test.cpp b/src/mongo/db/query/query_planner_array_test.cpp
index d4f53738164..3d2c3c1aa25 100644
--- a/src/mongo/db/query/query_planner_array_test.cpp
+++ b/src/mongo/db/query/query_planner_array_test.cpp
@@ -2076,4 +2076,47 @@ TEST_F(QueryPlannerTest, TypeArrayUsingStringAliasMustFetchAndFilter) {
         "bounds: {a: [['MinKey', 'MaxKey', true, true]]}}}}}");
 }
 
+TEST_F(QueryPlannerTest, CantExplodeMultikeyIxscanForSort) {
+    params.options &= ~QueryPlannerParams::INCLUDE_COLLSCAN;
+    const bool multikey = true;
+    addIndex(BSON("a" << 1 << "b" << 1), multikey);
+
+    runQueryAsCommand(fromjson("{find: 'testns', filter: {a: {$in: [1, 2]}}, sort: {b: 1}}"));
+
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: {sortKeyGen: {node: "
+        "{fetch: {filter: null, node: {ixscan: {pattern: {a: 1, b: 1}, filter: null, bounds: {a: "
+        "[[1,1,true,true], [2,2,true,true]], b: [['MinKey','MaxKey',true,true]]}}}}}}}}}");
+}
+
+TEST_F(QueryPlannerTest, CantExplodeMultikeyIxscanForSortWithPathLevelMultikeyMetadata) {
+    params.options &= ~QueryPlannerParams::INCLUDE_COLLSCAN;
+    MultikeyPaths multikeyPaths{std::set<size_t>{}, {0U}};
+    addIndex(BSON("a" << 1 << "b.c" << 1), multikeyPaths);
+
+    runQueryAsCommand(fromjson("{find: 'testns', filter: {a: {$in: [1, 2]}}, sort: {'b.c': 1}}"));
+
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{sort: {pattern: {'b.c': 1}, limit: 0, node: {sortKeyGen: {node: "
+        "{fetch: {filter: null, node: {ixscan: {pattern: {a: 1, 'b.c': 1}, filter: null, bounds: "
+        "{a: [[1,1,true,true], [2,2,true,true]], 'b.c': [['MinKey','MaxKey',true,true]]}}}}}}}}}");
+}
+
+TEST_F(QueryPlannerTest, CanExplodeMultikeyIndexScanForSortWhenSortFieldsAreNotMultikey) {
+    params.options &= ~QueryPlannerParams::INCLUDE_COLLSCAN;
+    MultikeyPaths multikeyPaths{{0U}, std::set<size_t>{}};
+    addIndex(BSON("a" << 1 << "b.c" << 1), multikeyPaths);
+
+    runQueryAsCommand(fromjson("{find: 'testns', filter: {a: {$in: [1, 2]}}, sort: {'b.c': 1}}"));
+
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: null, node: {mergeSort: {nodes: ["
+        "{ixscan: {pattern: {a: 1, 'b.c': 1}, filter: null,"
+        "bounds: {a: [[1,1,true,true]], 'b.c': [['MinKey','MaxKey',true,true]]}}},"
+        "{ixscan: {pattern: {a: 1, 'b.c': 1}, filter: null,"
+        "bounds: {a: [[2,2,true,true]], 'b.c': [['MinKey','MaxKey',true,true]]}}}]}}}}");
+}
 }  // namespace
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index e9095c01ceb..50208c7f7b1 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -758,7 +758,7 @@ void IndexScanNode::computeProperties() {
     // We cannot provide a sort which involves a multikey field. Prune such sort orders, if the
     // index is multikey.
     if (index.multikey) {
-        auto multikeyFields = getMultikeyFields(index.keyPattern, index.multikeyPaths);
+        multikeyFields = getMultikeyFields(index.keyPattern, index.multikeyPaths);
         for (auto sortsIt = _sorts.begin(); sortsIt != _sorts.end();) {
             bool foundMultikeyField = false;
             for (auto&& elt : *sortsIt) {
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index 9d3735edc02..e6b2e148d7c 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -495,6 +495,12 @@ struct IndexScanNode : public QuerySolutionNode {
     IndexBounds bounds;
 
     const CollatorInterface* queryCollator;
+
+    // The set of paths in the index key pattern which have at least one multikey path component, or
+    // empty if the index either is not multikey or does not have path-level multikeyness metadata.
+    //
+    // The correct set of paths is computed and stored here by computeProperties().
+    std::set<StringData> multikeyFields;
 };
 
 struct ProjectionNode : public QuerySolutionNode {
diff --git a/src/mongo/db/query/query_solution_test.cpp b/src/mongo/db/query/query_solution_test.cpp
index 0e07d9669a6..3e548869cc8 100644
--- a/src/mongo/db/query/query_solution_test.cpp
+++ b/src/mongo/db/query/query_solution_test.cpp
@@ -877,6 +877,32 @@ TEST(QuerySolutionTest, SimpleRangeAllEqualExcludesFieldWithMultikeyComponent) {
     ASSERT(node.getSort().count(BSON("e" << 1)));
 }
 
+TEST(QuerySolutionTest, MultikeyFieldsEmptyWhenIndexIsNotMultikey) {
+    IndexScanNode node{IndexEntry(BSON("a.b" << 1 << "c.d" << 1))};
+    node.index.multikey = false;
+    node.index.multikeyPaths = MultikeyPaths{};
+    node.computeProperties();
+    ASSERT(node.multikeyFields.empty());
+}
+
+TEST(QuerySolutionTest, MultikeyFieldsEmptyWhenIndexHasNoMultikeynessMetadata) {
+    IndexScanNode node{IndexEntry(BSON("a.b" << 1 << "c.d" << 1))};
+    node.index.multikey = true;
+    node.index.multikeyPaths = MultikeyPaths{};
+    node.computeProperties();
+    ASSERT(node.multikeyFields.empty());
+}
+
+TEST(QuerySolutionTest, MultikeyFieldsChosenCorrectlyWhenIndexHasPathLevelMultikeyMetadata) {
+    IndexScanNode node{IndexEntry(BSON("a.b" << 1 << "c.d" << 1 << "e.f" << 1))};
+    node.index.multikey = true;
+    node.index.multikeyPaths = MultikeyPaths{{0U}, {}, {0U, 1U}};
+    node.computeProperties();
+    ASSERT_EQ(node.multikeyFields.size(), 2U);
+    ASSERT(node.multikeyFields.count("a.b"));
+    ASSERT(node.multikeyFields.count("e.f"));
+}
+
 TEST(QuerySolutionTest, NonSimpleRangeAllEqualExcludesFieldWithMultikeyComponent) {
     IndexScanNode node{
         IndexEntry(BSON("a" << 1 << "b" << 1 << "c.z" << 1 << "d" << 1 << "e" << 1))};