SERVER-23618 Add collation support for partial indexes

3 files changed, 61 insertions, 8 deletions

diff --git a/src/mongo/db/matcher/SConscript b/src/mongo/db/matcher/SConscript
index c22650d40ca..ff82d538cba 100644
--- a/src/mongo/db/matcher/SConscript
+++ b/src/mongo/db/matcher/SConscript
@@ -99,6 +99,7 @@ env.CppUnitTest(
         'expression_algo_test.cpp',
     ],
     LIBDEPS=[
+        '$BUILD_DIR/mongo/db/query/collation/collator_interface_mock',
         'expression_algo',
     ],
 )
diff --git a/src/mongo/db/matcher/expression_algo.cpp b/src/mongo/db/matcher/expression_algo.cpp
index 02d25a5df9a..966c727501f 100644
--- a/src/mongo/db/matcher/expression_algo.cpp
+++ b/src/mongo/db/matcher/expression_algo.cpp
@@ -37,6 +37,7 @@
 #include "mongo/db/matcher/expression_leaf.h"
 #include "mongo/db/matcher/expression_tree.h"
 #include "mongo/db/pipeline/dependencies.h"
+#include "mongo/db/query/collation/collator_interface.h"
 
 namespace mongo {
 
@@ -94,7 +95,15 @@ bool _isSubsetOf(const ComparisonMatchExpression* lhs, const ComparisonMatchExpr
         return false;
     }
 
-    int cmp = compareElementValues(lhsData, rhsData);
+    // Either collator may be used by compareElementValues() here. If lhs (the query) contains
+    // string comparison, then _isSubsetOf will only be called if lhs and rhs have the same
+    // collator. Otherwise, the collator will not be used by compareElementValues().
+    if (!CollatorInterface::collatorsMatch(lhs->getCollator(), rhs->getCollator())) {
+        // TODO SERVER-23172: Check that lhsData does not contain string comparison in nested
+        // objects or arrays.
+        invariant(lhsData.type() != BSONType::String);
+    }
+    int cmp = compareElementValues(lhsData, rhsData, rhs->getCollator());
 
     // Check whether the two expressions are equivalent.
     if (lhs->matchType() == rhs->matchType() && cmp == 0) {
@@ -155,8 +164,7 @@ bool _isSubsetOf(const MatchExpression* lhs, const ComparisonMatchExpression* rh
         }
         for (BSONElement elem : ime->getEqualities()) {
             // Each element in the $in-array represents an equality predicate.
-            // TODO SERVER-23618: pass the appropriate collator to EqualityMatchExpression().
-            EqualityMatchExpression equality(nullptr);
+            EqualityMatchExpression equality(ime->getCollator());
             equality.init(lhs->path(), elem);
             if (!_isSubsetOf(&equality, rhs)) {
                 return false;
diff --git a/src/mongo/db/matcher/expression_algo_test.cpp b/src/mongo/db/matcher/expression_algo_test.cpp
index 95e7bcd62fc..ead421bc76e 100644
--- a/src/mongo/db/matcher/expression_algo_test.cpp
+++ b/src/mongo/db/matcher/expression_algo_test.cpp
@@ -39,6 +39,7 @@
 #include "mongo/db/matcher/expression_parser.h"
 #include "mongo/db/matcher/extensions_callback_disallow_extensions.h"
 #include "mongo/db/matcher/extensions_callback_noop.h"
+#include "mongo/db/query/collation/collator_interface_mock.h"
 #include "mongo/platform/decimal128.h"
 
 namespace mongo {
@@ -51,8 +52,8 @@ using std::unique_ptr;
  */
 class ParsedMatchExpression {
 public:
-    ParsedMatchExpression(const std::string& str) : _obj(fromjson(str)) {
-        const CollatorInterface* collator = nullptr;
+    ParsedMatchExpression(const std::string& str, const CollatorInterface* collator = nullptr)
+        : _obj(fromjson(str)) {
         StatusWithMatchExpression result =
             MatchExpressionParser::parse(_obj, ExtensionsCallbackDisallowExtensions(), collator);
         ASSERT_OK(result.getStatus());
@@ -489,9 +490,6 @@ TEST(ExpressionAlgoIsSubsetOf, RegexAndIn) {
     ParsedMatchExpression inRegexAOrEq1("{x: {$in: [/a/, 1]}}");
     ParsedMatchExpression inRegexAOrNull("{x: {$in: [/a/, null]}}");
 
-    ASSERT_TRUE(expression::isSubsetOf(inRegexA.get(), inRegexA.get()));
-    ASSERT_FALSE(expression::isSubsetOf(inRegexAbc.get(), inRegexA.get()));
-    ASSERT_FALSE(expression::isSubsetOf(inRegexA.get(), inRegexAOrEq1.get()));
     ASSERT_FALSE(expression::isSubsetOf(inRegexAOrEq1.get(), eq1.get()));
     ASSERT_FALSE(expression::isSubsetOf(inRegexA.get(), eqA.get()));
     ASSERT_FALSE(expression::isSubsetOf(inRegexAOrNull.get(), eqA.get()));
@@ -652,6 +650,52 @@ TEST(ExpressionAlgoIsSubsetOf, Compare_Exists_NE) {
     ASSERT_TRUE(expression::isSubsetOf(aNotEqualNull.get(), aExists.get()));
 }
 
+TEST(ExpressionAlgoIsSubsetOf, CollationAwareStringComparison) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    ParsedMatchExpression lhs("{a: {$gt: 'abc'}}", &collator);
+    ParsedMatchExpression rhs("{a: {$gt: 'cba'}}", &collator);
+
+    ASSERT_TRUE(expression::isSubsetOf(lhs.get(), rhs.get()));
+
+    ParsedMatchExpression lhsLT("{a: {$lt: 'abc'}}", &collator);
+    ParsedMatchExpression rhsLT("{a: {$lt: 'cba'}}", &collator);
+
+    ASSERT_FALSE(expression::isSubsetOf(lhsLT.get(), rhsLT.get()));
+}
+
+TEST(ExpressionAlgoIsSubsetOf, CollationAwareStringComparisonIn) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    ParsedMatchExpression lhsAllGTcba("{a: {$in: ['abc', 'cbc']}}", &collator);
+    ParsedMatchExpression lhsSomeGTcba("{a: {$in: ['abc', 'aba']}}", &collator);
+    ParsedMatchExpression rhs("{a: {$gt: 'cba'}}", &collator);
+
+    ASSERT_TRUE(expression::isSubsetOf(lhsAllGTcba.get(), rhs.get()));
+    ASSERT_FALSE(expression::isSubsetOf(lhsSomeGTcba.get(), rhs.get()));
+
+    ParsedMatchExpression rhsLT("{a: {$lt: 'cba'}}", &collator);
+
+    ASSERT_FALSE(expression::isSubsetOf(lhsAllGTcba.get(), rhsLT.get()));
+    ASSERT_FALSE(expression::isSubsetOf(lhsSomeGTcba.get(), rhsLT.get()));
+}
+
+TEST(ExpressionAlgoIsSubsetOf, NonMatchingCollationsNoStringComparisonLHS) {
+    CollatorInterfaceMock collatorAlwaysEqual(CollatorInterfaceMock::MockType::kAlwaysEqual);
+    CollatorInterfaceMock collatorReverseString(CollatorInterfaceMock::MockType::kReverseString);
+    ParsedMatchExpression lhs("{a: {b: 1}}", &collatorAlwaysEqual);
+    ParsedMatchExpression rhs("{a: {$lt: {b: 'abc'}}}", &collatorReverseString);
+
+    ASSERT_TRUE(expression::isSubsetOf(lhs.get(), rhs.get()));
+}
+
+TEST(ExpressionAlgoIsSubsetOf, NonMatchingCollationsNoStringComparison) {
+    CollatorInterfaceMock collatorAlwaysEqual(CollatorInterfaceMock::MockType::kAlwaysEqual);
+    CollatorInterfaceMock collatorReverseString(CollatorInterfaceMock::MockType::kReverseString);
+    ParsedMatchExpression lhs("{a: 1}", &collatorAlwaysEqual);
+    ParsedMatchExpression rhs("{a: {$gt: 0}}", &collatorReverseString);
+
+    ASSERT_TRUE(expression::isSubsetOf(lhs.get(), rhs.get()));
+}
+
 TEST(IsIndependent, AndIsIndependentOnlyIfChildrenAre) {
     BSONObj matchPredicate = fromjson("{$and: [{a: 1}, {b: 1}]}");
     const CollatorInterface* collator = nullptr;