SERVER-48993 explodeForSort can produce incorrect query plan

(cherry picked from commit 7dd37622622362d0ef689b91de565c8c053c838e)

5 files changed, 379 insertions, 0 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index f52fd671078..e70f05258e7 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -35,6 +35,8 @@ replica_sets_jscore_multiversion_passthrough:
   test_file: jstests/core/geo_near_tailable.js
 - ticket: SERVER-48614
   test_file: jstests/core/wildcard_index_partial_index.js
+- ticket: SERVER-48993
+  test_file: jstests/core/explode_for_sort_collation.js
 
 replica_sets_multiversion:
 - ticket: SERVER-42825
@@ -69,15 +71,21 @@ sharding_multiversion:
   test_file: jstests/sharding/explain_cmd.js
 - ticket: SERVER-43889
   test_file: jstests/sharding/retryable_writes.js
+- ticket: SERVER-48993
+  test_file: jstests/core/explode_for_sort_collation.js
 
 sharding_jscore_multiversion_passthrough:
 - ticket: SERVER-47773
   test_file: jstests/core/geo_near_tailable.js
 - ticket: SERVER-48614
   test_file: jstests/core/wildcard_index_partial_index.js
+- ticket: SERVER-48993
+  test_file: jstests/core/explode_for_sort_collation.js
 
 sharded_collections_jscore_multiversion_passthrough:
 - ticket: SERVER-47773
   test_file: jstests/core/geo_near_tailable.js
 - ticket: SERVER-48614
   test_file: jstests/core/wildcard_index_partial_index.js
+- ticket: SERVER-48993
+  test_file: jstests/core/explode_for_sort_collation.js
diff --git a/jstests/core/explode_for_sort_collation.js b/jstests/core/explode_for_sort_collation.js
new file mode 100644
index 00000000000..fcdb3629bce
--- /dev/null
+++ b/jstests/core/explode_for_sort_collation.js
@@ -0,0 +1,203 @@
+/**
+ * Tests explode for sort query planner behavior with collated queries and indexes. This is a test
+ * for SERVER-48993.
+ * @tags: [requires_find_command]
+ */
+(function() {
+"use strict";
+
+const testDB = db.getSiblingDB(jsTestName());
+
+// Drop the test database.
+assert.commandWorked(testDB.dropDatabase());
+
+const coll = testDB.explode_for_sort;
+
+// Executes a test case that creates an index, inserts documents, issues a find query on a
+// collection and compares the results with the expected collection.
+function executeQueryTestCase(testCase) {
+    jsTestLog(tojson(testCase));
+
+    // Create a collection.
+    coll.drop();
+    assert.commandWorked(testDB.createCollection(coll.getName()));
+
+    // Create an index.
+    assert.commandWorked(coll.createIndex(testCase.indexKeyPattern, testCase.indexOptions));
+
+    // Insert some documents into the collection.
+    assert.commandWorked(coll.insert(testCase.inputDocuments));
+
+    // Run a find query with optionally specified collation.
+    let cursor = coll.find(testCase.filter).sort(testCase.sort);
+    if (testCase.findCollation !== undefined) {
+        cursor = cursor.collation(testCase.findCollation);
+    }
+    const actualResults = cursor.toArray();
+
+    // Compare results to expected.
+    assert.eq(actualResults, testCase.expectedResults);
+}
+
+const standardInputDocuments =
+    [{_id: 0, a: 0, b: "CC"}, {_id: 1, a: 0, b: "AA"}, {_id: 2, a: 0, b: "bb"}];
+
+const testCases = [
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key together with a
+        // sort on a suffix of the index key returns correct results when the index is a compound
+        // index with a non-simple collation and the query does not have an explicit collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 1}},
+        filter: {a: 0},
+        sort: {b: 1},
+        inputDocuments: standardInputDocuments,
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 0, a: 0, b: "CC"}, {_id: 2, a: 0, b: "bb"}]
+    },
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key together with a
+        // sort on a suffix of the index key returns correct results when the index is a compound
+        // index with a non-simple collation  and the query explicitly specifies the simple
+        // collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 1}},
+        filter: {a: 0},
+        sort: {b: 1},
+        findCollation: {locale: "simple"},
+        inputDocuments: standardInputDocuments,
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 0, a: 0, b: "CC"}, {_id: 2, a: 0, b: "bb"}]
+    },
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key together with a
+        // sort on a suffix of the index key returns correct results when the index is a compound
+        // index with a simple collation and the query explicitly specifies a non-simple collation.
+        indexKeyPattern: {a: 1, b: 1},
+        filter: {a: 0},
+        sort: {b: 1},
+        findCollation: {locale: "en_US", strength: 1},
+        inputDocuments: standardInputDocuments,
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 2, a: 0, b: "bb"}, {_id: 0, a: 0, b: "CC"}]
+    },
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key together with a
+        // sort on a suffix of the index key returns correct results when the index is a compound
+        // index with a simple collation and the query explicitly specifies a non-simple collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "simple"}},
+        filter: {a: 0},
+        sort: {b: 1},
+        findCollation: {locale: "en_US", strength: 1},
+        inputDocuments: standardInputDocuments,
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 2, a: 0, b: "bb"}, {_id: 0, a: 0, b: "CC"}]
+    },
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key together with a
+        // sort on a suffix of the index key returns correct results when the index is a compound
+        // index with a non-simple collation that is different from the query's.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 5}},
+        filter: {a: 0},
+        sort: {b: 1},
+        findCollation: {locale: "en_US", strength: 1},
+        inputDocuments: standardInputDocuments,
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 2, a: 0, b: "bb"}, {_id: 0, a: 0, b: "CC"}]
+    },
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key, a collatable
+        // range-query on the suffix, and a sort on the suffix of the index key returns correct
+        // results when the index is a compound index with a non-simple collation and the query does
+        // not have an explicit collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 1}},
+        filter: {a: 0, b: {$gte: 'A', $lt: 'D'}},
+        sort: {b: 1},
+        inputDocuments: standardInputDocuments,
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 0, a: 0, b: "CC"}]
+    },
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key, a collatable
+        // range-query and a sort on a prefix of a suffix of the index key returns correct results
+        // when the index is a compound index with a non-simple collation and the query does not
+        // have an explicit collation.
+        indexKeyPattern: {a: 1, b: 1, c: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 1}},
+        filter: {a: 0, b: {$gte: 'A', $lt: 'D'}},
+        sort: {b: 1},
+        inputDocuments: standardInputDocuments,
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 0, a: 0, b: "CC"}]
+    },
+    {
+        // Verifies that a non-collatable multi-point query on the prefix of the index key, a
+        // collatable range-query on the suffix, and a sort on the suffix of the index key returns
+        // correct results when the index is a compound index with a non-simple collation and the
+        // query does not have an explicit collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 1}},
+        filter: {a: {$in: [0, 2]}, b: {$gte: 'A', $lt: 'D'}},
+        sort: {b: 1},
+        inputDocuments: [
+            {_id: 0, a: 0, b: "CC"},
+            {_id: 1, a: 0, b: "AA"},
+            {_id: 2, a: 0, b: "bb"},
+            {_id: 3, a: 2, b: "BB"}
+        ],
+        expectedResults: [{_id: 1, a: 0, b: "AA"}, {_id: 3, a: 2, b: "BB"}, {_id: 0, a: 0, b: "CC"}]
+    },
+    {
+        // Verifies that a non-collatable multi-point query on the prefix of the index key, a
+        // non-collatable range-query on the suffix, and a sort on the suffix of the index key
+        // returns correct results when the index is a compound index with a non-simple collation
+        // and the query does not have an explicit collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 1}},
+        filter: {a: {$in: [0, 2]}, b: {$gte: 0, $lt: 10}},
+        sort: {b: 1},
+        inputDocuments: [
+            {_id: 0, a: 0, b: 6},
+            {_id: 1, a: 0, b: 10},
+            {_id: 2, a: 0, b: "bb"},
+            {_id: 3, a: 2, b: 5},
+            {_id: 4, a: 2, b: 4}
+        ],
+        expectedResults: [{_id: 4, a: 2, b: 4}, {_id: 3, a: 2, b: 5}, {_id: 0, a: 0, b: 6}]
+    },
+    {
+        // Verifies that a non-collatable multi-point query on the prefix of the index key, a
+        // non-collatable range-query on the suffix, and a sort on the suffix of the index key
+        // returns correct results when the index is a compound index with a simple collation
+        // and the query explicitly specifies a non-simple collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "simple"}},
+        filter: {a: {$in: [0, 2]}, b: {$gte: 0, $lt: 10}},
+        sort: {b: 1},
+        findCollation: {locale: "en_US", strength: 1},
+        inputDocuments: [
+            {_id: 0, a: 0, b: 6},
+            {_id: 1, a: 0, b: 10},
+            {_id: 2, a: 0, b: "bb"},
+            {_id: 3, a: 2, b: 5},
+            {_id: 4, a: 2, b: 4}
+        ],
+        expectedResults: [{_id: 4, a: 2, b: 4}, {_id: 3, a: 2, b: 5}, {_id: 0, a: 0, b: 6}]
+    },
+    {
+        // Verifies that a non-collatable point-query on the prefix of the index key, a
+        // non-collatable range-query on the suffix, and a sort on the suffix of the index key
+        // returns correct results when the index is a compound index with a non-simple collation
+        // and the query does not have an explicit collation.
+        indexKeyPattern: {a: 1, b: 1},
+        indexOptions: {collation: {locale: "en_US", strength: 1}},
+        filter: {a: 0, b: {$gte: 0, $lt: 10}},
+        sort: {b: 1},
+        inputDocuments: [
+            {_id: 0, a: 0, b: 6},
+            {_id: 1, a: 0, b: 5},
+            {_id: 2, a: 0, b: "bb"},
+            {_id: 3, a: 0, b: 4}
+        ],
+        expectedResults: [{_id: 3, a: 0, b: 4}, {_id: 1, a: 0, b: 5}, {_id: 0, a: 0, b: 6}]
+    }
+];
+
+testCases.forEach(executeQueryTestCase);
+}());
\ No newline at end of file
diff --git a/jstests/libs/parallelTester.js b/jstests/libs/parallelTester.js
index 2a2a1576763..3b9f0c6809e 100644
--- a/jstests/libs/parallelTester.js
+++ b/jstests/libs/parallelTester.js
@@ -212,6 +212,7 @@ if (typeof _threadInject != "undefined") {
         // The following tests cannot run when shell readMode is legacy.
         if (db.getMongo().readMode() === "legacy") {
             var requires_find_command = [
+                "explode_for_sort_collation.js",
                 "update_pipeline_shell_helpers.js",
                 "update_with_pipeline.js",
                 "views/views_aggregation.js",
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 3c4b601d1aa..ab8698c4e43 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -542,6 +542,20 @@ bool QueryPlannerAnalysis::explodeForSort(const CanonicalQuery& query,
             }
         }
 
+        // An index whose collation does not match the query's cannot provide a sort if sort-by
+        // fields can contain collatable values.
+        if (!CollatorInterface::collatorsMatch(isn->index.collator, query.getCollator())) {
+            auto fieldsWithStringBounds =
+                IndexScanNode::getFieldsWithStringBounds(bounds, isn->index.keyPattern);
+            for (auto&& element : desiredSort) {
+                if (fieldsWithStringBounds.count(element.fieldNameStringData()) > 0) {
+                    // The field can contain collatable values and therefore we cannot use the index
+                    // to provide the sort.
+                    return false;
+                }
+            }
+        }
+
         // Do some bookkeeping to see how many ixscans we'll create total.
         totalNumScans += numScans;
 
diff --git a/src/mongo/db/query/query_planner_collation_test.cpp b/src/mongo/db/query/query_planner_collation_test.cpp
index 5cefee957c1..aec22d66680 100644
--- a/src/mongo/db/query/query_planner_collation_test.cpp
+++ b/src/mongo/db/query/query_planner_collation_test.cpp
@@ -612,6 +612,159 @@ TEST_F(QueryPlannerTest, CannotUseIndexWhenQueryHasDifferentNonSimpleCollationTh
         "{sortKeyGen:{node: {cscan: {dir: 1}}}}}}}}");
 }
 
+// This test verifies that an in-memory sort stage is added and sort provided by an index is not
+// used when the collection has a compound index with a non-simple collation and we issue a
+// non-collatable point-query on the prefix of the index key together with a sort on a suffix of the
+// index key. This is a test for SERVER-48993.
+TEST_F(QueryPlannerTest,
+       MustSortInMemoryWhenPointPrefixQueryHasSimpleCollationButIndexHasNonSimpleCollation) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kToLowerString);
+    addIndex(fromjson("{a: 1, b: 1}"), &collator);
+
+    // No explicit collation on the query. This will implicitly use the simple collation since the
+    // collection does not have a default collation.
+    runQueryAsCommand(fromjson("{find: 'testns', filter: {a: 2}, sort: {b: 1}}"));
+
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{cscan: {dir: 1}}}}}}");
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{fetch: {node: "
+        "{ixscan: {pattern: {a: 1, b: 1}}}}}}}}}");
+
+    // A query with an explicit simple collation.
+    runQueryAsCommand(
+        fromjson("{find: 'testns', filter: {a: 2}, sort: {b: 1}, collation: {locale: 'simple'}}"));
+
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{cscan: {dir: 1}}}}}}");
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{fetch: {node: "
+        "{ixscan: {pattern: {a: 1, b: 1}}}}}}}}}");
+}
+
+// This test verifies that an in-memory sort stage is added and sort provided by an index is not
+// used when the collection has a compound index with a non-specified collation and we issue a
+// non-collatable point-query on the prefix of the index key together with a sort on the suffix and
+// an explicit non-simple collation. This is a test for SERVER-48993.
+TEST_F(QueryPlannerTest,
+       MustSortInMemoryWhenPointPrefixQueryHasNonSimpleCollationButIndexHasSimpleCollation) {
+    addIndex(fromjson("{a: 1, b: 1}"));
+
+    runQueryAsCommand(
+        fromjson("{find: 'testns', filter: {a: 2}, sort: {b: 1}, collation: {locale: "
+                 "'reverse'}}"));
+
+    assertSolutionExists(
+        "{fetch: {node: "
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{ixscan: {pattern: {a: 1, b: 1}}}}}}}}}");
+}
+
+// This test verifies that an in-memory sort stage is added and sort provided by indexes is not used
+// when the collection has a compound index with a non-simple collation and we issue a
+// non-collatable point-query on the prefix of the index key together with a sort on the suffix and
+// an explicit non-simple collation that differs from the index collation. This is a test for
+// SERVER-48993.
+TEST_F(QueryPlannerTest, MustSortInMemoryWhenPointPrefixQueryCollationDoesNotMatchIndexCollation) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kToLowerString);
+    addIndex(fromjson("{a: 1, b: 1}"), &collator);
+
+    runQueryAsCommand(
+        fromjson("{find: 'testns', filter: {a: 2}, sort: {b: 1}, collation: {locale: 'reverse'}}"));
+
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{cscan: {dir: 1}}}}}}");
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{fetch: {node: "
+        "{ixscan: {pattern: {a: 1, b: 1}}}}}}}}}");
+}
+
+// This test verifies that a sort is provided by an index when the collection has a compound index
+// with a non-simple collation and we issue a query with a different non-simple collation is a
+// non-collatable point-query on the prefix, a non-collatable range-query on the suffix, and a sort
+// on the suffix key. This is a test for SERVER-48993.
+TEST_F(QueryPlannerTest,
+       IndexCanSortWhenPointPrefixQueryCollationDoesNotMatchIndexButSortRangeIsNonCollatable) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kToLowerString);
+    addIndex(fromjson("{a: 1, b: 1}"), &collator);
+
+    runQueryAsCommand(
+        fromjson("{find: 'testns', filter: {a: 2, b: {$gte: 0, $lt: 10}}, sort: {b: 1}}"));
+
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{cscan: {dir: 1}}}}}}");
+    assertSolutionExists(
+        "{fetch: {node: "
+        "{ixscan: {pattern: {a: 1, b: 1}, bounds: {a: [[2, 2, true, true]], b: [[0, 10, true, "
+        "false]]}}}}}");
+}
+
+// This test verifies that an in-memory sort stage is added when the collection has a compound index
+// with a non-simple collation and we issue a query with a different non-simple collation is a
+// non-collatable point-query on the prefix, a collatable range-query on the suffix, and a sort on
+// the suffix key. This is a test for SERVER-48993.
+TEST_F(QueryPlannerTest,
+       MustSortInMemoryWhenPointPrefixQueryCollationDoesNotMatchIndexAndSortRangeIsCollatable) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kToLowerString);
+    addIndex(fromjson("{a: 1, b: 1}"), &collator);
+
+    runQueryAsCommand(
+        fromjson("{find: 'testns', filter: {a: 2, b: {$gte: 'B', $lt: 'T'}}, sort: {b: 1}}"));
+
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{cscan: {dir: 1}}}}}}");
+    assertSolutionExists(
+        "{sort: {pattern: {b: 1}, limit: 0, node: "
+        "{sortKeyGen: {node:"
+        "{fetch: {node: "
+        "{ixscan: {pattern: {a: 1, b: 1}, bounds: {a: [[2, 2, true, true]]}}}}}}}}}");
+}
+
+// This test verifies that a SORT_MERGE stage is added when the collection has a compound index with
+// a non-simple collation and we issue a query with a different non-simple collation is a
+// non-collatable multi-point query on the prefix, a non-collatable range-query on the suffix, and a
+// sort on the suffix key.This is a test for SERVER-48993.
+TEST_F(QueryPlannerTest,
+       CanExplodeForSortWhenPointPrefixQueryCollationDoesNotMatchIndexButSortRangeIsNonCollatable) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kToLowerString);
+    addIndex(fromjson("{a: 1, b: 1, c: 1}"), &collator);
+
+    runQueryAsCommand(fromjson(
+        "{find: 'testns', filter: {a: {$in: [2, 5]}, b: {$gte: 0, $lt: 10}}, sort: {b: 1}}"));
+
+    assertNumSolutions(2U);
+    assertSolutionExists(
+        "{fetch: {node: "
+        "{mergeSort: {nodes: {"
+        "n0: {ixscan: {pattern: {a: 1, b: 1, c: 1}, bounds: {a: [[2, 2, true, true]], b: [[0, 10, "
+        "true, false]]}}}, "
+        "n1: {ixscan: {pattern: {a: 1, b: 1, c: 1}, bounds: {a: [[5, 5, true, true]], b: [[0, 10, "
+        "true, false]]}}} "
+        "}}}}}");
+}
+
 /**
  * This test confirms that we place a fetch stage before sortKeyGen in the case where both query
  * and index have the same non-simple collation. To handle this scenario without this fetch would