SERVER-23618 Add collation support for partial indexes

4 files changed, 72 insertions, 9 deletions

diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp
index 2b76ac930f9..faad7ca6aae 100644
--- a/src/mongo/db/query/query_planner.cpp
+++ b/src/mongo/db/query/query_planner.cpp
@@ -851,12 +851,6 @@ Status QueryPlanner::plan(const CanonicalQuery& query,
                     continue;
                 }
 
-                // If the index collation differs from the query collation, the index should not be
-                // used to provide a sort, because strings will be ordered incorrectly.
-                if (!CollatorInterface::collatorsMatch(index.collator, query.getCollator())) {
-                    continue;
-                }
-
                 // Partial indexes can only be used to provide a sort only if the query predicate is
                 // compatible.
                 if (index.filterExpr && !expression::isSubsetOf(query.root(), index.filterExpr)) {
diff --git a/src/mongo/db/query/query_planner_partialidx_test.cpp b/src/mongo/db/query/query_planner_partialidx_test.cpp
index 6f27939b10a..476e09041e3 100644
--- a/src/mongo/db/query/query_planner_partialidx_test.cpp
+++ b/src/mongo/db/query/query_planner_partialidx_test.cpp
@@ -28,6 +28,7 @@
 
 #include "mongo/platform/basic.h"
 
+#include "mongo/db/query/collation/collator_interface_mock.h"
 #include "mongo/db/query/query_planner.h"
 #include "mongo/db/query/query_planner_test_fixture.h"
 
@@ -443,5 +444,55 @@ TEST_F(QueryPlannerTest, PartialIndexNor) {
     assertNumSolutions(0U);
 }
 
+TEST_F(QueryPlannerTest, PartialIndexStringComparisonMatchingCollators) {
+    params.options = QueryPlannerParams::NO_TABLE_SCAN;
+    BSONObj filterObj(fromjson("{a: {$gt: 'cba'}}"));
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    std::unique_ptr<MatchExpression> filterExpr = parseMatchExpression(filterObj, &collator);
+    addIndex(fromjson("{a: 1}"), filterExpr.get(), &collator);
+
+    runQueryAsCommand(
+        fromjson("{find: 'testns', filter: {a: 'abc'}, collation: {locale: 'reverse'}}"));
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: {a: 'abc'}, collation: {locale: 'reverse'}, node: {ixscan: "
+        "{filter: null, pattern: {a: 1}, "
+        "bounds: {a: [['cba', 'cba', true, true]]}}}}}");
+
+    runQueryAsCommand(
+        fromjson("{find: 'testns', filter: {a: 'zaa'}, collation: {locale: 'reverse'}}"));
+    assertNumSolutions(0U);
+}
+
+TEST_F(QueryPlannerTest, PartialIndexFilterStringComparisonNonMatchingCollators) {
+    params.options = QueryPlannerParams::NO_TABLE_SCAN;
+    BSONObj filterObj(fromjson("{a: {$lt: {b: 'abc'}}}"));
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    std::unique_ptr<MatchExpression> filterExpr = parseMatchExpression(filterObj, &collator);
+    addIndex(fromjson("{a: 1}"), filterExpr.get(), &collator);
+
+    runQuery(fromjson("{a: {b: 1}}"));
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: null, node: {ixscan: "
+        "{filter: null, pattern: {a: 1}, "
+        "bounds: {a: [[{b: 1}, {b: 1}, true, true]]}}}}}");
+}
+
+TEST_F(QueryPlannerTest, PartialIndexNoStringComparisonNonMatchingCollators) {
+    params.options = QueryPlannerParams::NO_TABLE_SCAN;
+    BSONObj filterObj(fromjson("{a: {$gt: 0}}"));
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kAlwaysEqual);
+    std::unique_ptr<MatchExpression> filterExpr = parseMatchExpression(filterObj, &collator);
+    addIndex(fromjson("{a: 1}"), filterExpr.get(), &collator);
+
+    runQueryAsCommand(fromjson("{find: 'testns', filter: {a: 1}, collation: {locale: 'reverse'}}"));
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{fetch: {filter: null, node: {ixscan: "
+        "{filter: null, pattern: {a: 1}, "
+        "bounds: {a: [[1, 1, true, true]]}}}}}");
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_planner_test_fixture.cpp b/src/mongo/db/query/query_planner_test_fixture.cpp
index aae4d95372f..c49df6ab106 100644
--- a/src/mongo/db/query/query_planner_test_fixture.cpp
+++ b/src/mongo/db/query/query_planner_test_fixture.cpp
@@ -139,6 +139,19 @@ void QueryPlannerTest::addIndex(BSONObj keyPattern, const CollatorInterface* col
     params.indices.push_back(entry);
 }
 
+void QueryPlannerTest::addIndex(BSONObj keyPattern,
+                                MatchExpression* filterExpr,
+                                const CollatorInterface* collator) {
+    const bool sparse = false;
+    const bool unique = false;
+    const bool multikey = false;
+    const char name[] = "my_partial_index_with_collator";
+    const BSONObj infoObj;
+    IndexEntry entry(keyPattern, multikey, sparse, unique, name, filterExpr, infoObj);
+    entry.collator = collator;
+    params.indices.push_back(entry);
+}
+
 void QueryPlannerTest::runQuery(BSONObj query) {
     runQuerySortProjSkipNToReturn(query, BSONObj(), BSONObj(), 0, 0);
 }
@@ -399,8 +412,8 @@ void QueryPlannerTest::assertHasOneSolutionOf(const std::vector<std::string>& so
     FAIL(ss);
 }
 
-std::unique_ptr<MatchExpression> QueryPlannerTest::parseMatchExpression(const BSONObj& obj) {
-    const CollatorInterface* collator = nullptr;
+std::unique_ptr<MatchExpression> QueryPlannerTest::parseMatchExpression(
+    const BSONObj& obj, const CollatorInterface* collator) {
     StatusWithMatchExpression status =
         MatchExpressionParser::parse(obj, ExtensionsCallbackDisallowExtensions(), collator);
     if (!status.isOK()) {
diff --git a/src/mongo/db/query/query_planner_test_fixture.h b/src/mongo/db/query/query_planner_test_fixture.h
index 2e274274fb1..6c649702c74 100644
--- a/src/mongo/db/query/query_planner_test_fixture.h
+++ b/src/mongo/db/query/query_planner_test_fixture.h
@@ -68,6 +68,10 @@ protected:
 
     void addIndex(BSONObj keyPattern, const CollatorInterface* collator);
 
+    void addIndex(BSONObj keyPattern,
+                  MatchExpression* filterExpr,
+                  const CollatorInterface* collator);
+
     //
     // Execute planner.
     //
@@ -192,7 +196,8 @@ protected:
     /**
      * Helper function to parse a MatchExpression.
      */
-    static std::unique_ptr<MatchExpression> parseMatchExpression(const BSONObj& obj);
+    static std::unique_ptr<MatchExpression> parseMatchExpression(
+        const BSONObj& obj, const CollatorInterface* collator = nullptr);
 
     //
     // Data members.