SERVER-23172 Allow use of indices for collation-aware queries that match nested objects or arrays.

3 files changed, 159 insertions, 13 deletions

diff --git a/src/mongo/db/query/collation/collation_index_key.cpp b/src/mongo/db/query/collation/collation_index_key.cpp
index eb25acc3f79..ed0f348a887 100644
--- a/src/mongo/db/query/collation/collation_index_key.cpp
+++ b/src/mongo/db/query/collation/collation_index_key.cpp
@@ -33,26 +33,91 @@
 #include "mongo/bson/bsonobj.h"
 #include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/db/query/collation/collator_interface.h"
+#include "mongo/util/assert_util.h"
 
 namespace mongo {
 
-// TODO SERVER-23172: Update this to consider strings inside nested objects or arrays.
+void CollationIndexKey::translate(StringData fieldName,
+                                  BSONElement element,
+                                  const CollatorInterface* collator,
+                                  BSONObjBuilder* out) {
+    invariant(collator);
+
+    switch (element.type()) {
+        case BSONType::String: {
+            out->append(fieldName,
+                        collator->getComparisonKey(element.valueStringData()).getKeyData());
+            return;
+        }
+        case BSONType::Object: {
+            BSONObjBuilder bob(out->subobjStart(fieldName));
+            for (const BSONElement& elt : element.embeddedObject()) {
+                translate(elt.fieldNameStringData(), elt, collator, &bob);
+            }
+            bob.doneFast();
+            return;
+        }
+        case BSONType::Array: {
+            BSONArrayBuilder bab(out->subarrayStart(fieldName));
+            for (const BSONElement& elt : element.Array()) {
+                translate(elt, collator, &bab);
+            }
+            bab.doneFast();
+            return;
+        }
+        default:
+            out->appendAs(element, fieldName);
+    }
+}
+
+void CollationIndexKey::translate(BSONElement element,
+                                  const CollatorInterface* collator,
+                                  BSONArrayBuilder* out) {
+    invariant(collator);
+
+    switch (element.type()) {
+        case BSONType::String: {
+            out->append(collator->getComparisonKey(element.valueStringData()).getKeyData());
+            return;
+        }
+        case BSONType::Object: {
+            BSONObjBuilder bob(out->subobjStart());
+            for (const BSONElement& elt : element.embeddedObject()) {
+                translate(elt.fieldNameStringData(), elt, collator, &bob);
+            }
+            bob.doneFast();
+            return;
+        }
+        case BSONType::Array: {
+            BSONArrayBuilder bab(out->subarrayStart());
+            for (const BSONElement& elt : element.Array()) {
+                translate(elt, collator, &bab);
+            }
+            bab.doneFast();
+            return;
+        }
+        default:
+            out->append(element);
+    }
+}
+
+// TODO SERVER-24674: We may want to check that objects and arrays actually do contain strings
+// before returning true.
 bool CollationIndexKey::shouldUseCollationIndexKey(BSONElement elt,
                                                    const CollatorInterface* collator) {
-    return collator && elt.type() == BSONType::String;
+    return collator && isCollatableType(elt.type());
 }
 
-// TODO SERVER-23172: Update this to convert strings inside nested objects or arrays to their
-// corresponding comparison keys.
 void CollationIndexKey::collationAwareIndexKeyAppend(BSONElement elt,
                                                      const CollatorInterface* collator,
                                                      BSONObjBuilder* out) {
-    if (shouldUseCollationIndexKey(elt, collator)) {
-        auto comparisonKey = collator->getComparisonKey(elt.valueStringData());
-        out->append("", comparisonKey.getKeyData());
-    } else {
+    invariant(out);
+    if (!collator) {
         out->appendAs(elt, "");
+        return;
     }
+
+    translate("", elt, collator, out);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/collation/collation_index_key.h b/src/mongo/db/query/collation/collation_index_key.h
index 7d86935b11a..dd733e5a08c 100644
--- a/src/mongo/db/query/collation/collation_index_key.h
+++ b/src/mongo/db/query/collation/collation_index_key.h
@@ -28,6 +28,8 @@
 
 #pragma once
 
+#include "mongo/bson/bsontypes.h"
+
 namespace mongo {
 
 class BSONElement;
@@ -41,7 +43,17 @@ class CollatorInterface;
 class CollationIndexKey {
 public:
     /**
-     * Returns true if the index key for 'elt' should be a comparison key generated by 'collator'.
+     * Returns true if type is affected in comparison or sort order by collation.
+     */
+    static bool isCollatableType(BSONType type) {
+        // TODO SERVER-24674 We may want to replace calls to this to shouldUseCollationIndexKey.
+        return type == BSONType::String || type == BSONType::Object || type == BSONType::Array;
+    }
+
+    /**
+     * Returns true if the index key for 'elt' may contain a comparison key generated by 'collator'.
+     *
+     * This is the case when 'elt' is a String, Array, or Object.
      */
     static bool shouldUseCollationIndexKey(BSONElement elt, const CollatorInterface* collator);
 
@@ -52,6 +64,20 @@ public:
     static void collationAwareIndexKeyAppend(BSONElement elt,
                                              const CollatorInterface* collator,
                                              BSONObjBuilder* out);
+
+private:
+    // Translate all strings in 'element' into comparison keys using 'collator'. The result is
+    // appended to the BSONObjBuilder out, with field name fieldName.
+    static void translate(StringData fieldName,
+                          BSONElement element,
+                          const CollatorInterface* collator,
+                          BSONObjBuilder* out);
+
+    // Translate all strings in 'element' into comparison keys using 'collator'.  The result is
+    // append to the BSONArrayBuilder out.
+    static void translate(BSONElement element,
+                          const CollatorInterface* collator,
+                          BSONArrayBuilder* out);
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/collation/collation_index_key_test.cpp b/src/mongo/db/query/collation/collation_index_key_test.cpp
index 88d5b13c70c..bcd341c9488 100644
--- a/src/mongo/db/query/collation/collation_index_key_test.cpp
+++ b/src/mongo/db/query/collation/collation_index_key_test.cpp
@@ -31,6 +31,7 @@
 #include "mongo/db/query/collation/collation_index_key.h"
 
 #include "mongo/bson/bsonobjbuilder.h"
+#include "mongo/bson/json.h"
 #include "mongo/db/query/collation/collator_interface_mock.h"
 #include "mongo/unittest/unittest.h"
 
@@ -38,26 +39,40 @@ namespace {
 
 using namespace mongo;
 
-TEST(CollationIndexKeyTest, ShouldUseCollationKeyFalseWithNullCollator) {
+TEST(CollationIndexKeyTest, ShouldUseCollationIndexKeyFalseWithNullCollator) {
     BSONObj obj = BSON("foo"
                        << "string");
     ASSERT_FALSE(CollationIndexKey::shouldUseCollationIndexKey(obj.firstElement(), nullptr));
 }
 
-TEST(CollationIndexKeyTest, ShouldUseCollationKeyFalseWithNonStringElement) {
+TEST(CollationIndexKeyTest, ShouldUseCollationIndexKeyTrueWithObjectElement) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj obj = BSON("foo" << BSON("bar"
                                      << "string"));
-    ASSERT_FALSE(CollationIndexKey::shouldUseCollationIndexKey(obj.firstElement(), &collator));
+    ASSERT_TRUE(CollationIndexKey::shouldUseCollationIndexKey(obj.firstElement(), &collator));
 }
 
-TEST(CollationIndexKeyTest, ShouldUseCollationKeyTrueWithStringElement) {
+TEST(CollationIndexKeyTest, ShouldUseCollationIndexKeyTrueWithArrayElement) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj obj = BSON("foo" << BSON_ARRAY("one"
+                                           << "two"));
+    ASSERT_TRUE(CollationIndexKey::shouldUseCollationIndexKey(obj.firstElement(), &collator));
+}
+
+TEST(CollationIndexKeyTest, ShouldUseCollationIndexKeyTrueWithStringElement) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj obj = BSON("foo"
                        << "string");
     ASSERT_TRUE(CollationIndexKey::shouldUseCollationIndexKey(obj.firstElement(), &collator));
 }
 
+TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyAppendsElementWithNullCollator) {
+    BSONObj dataObj = BSON("test" << 1);
+    BSONObjBuilder out;
+    CollationIndexKey::collationAwareIndexKeyAppend(dataObj.firstElement(), nullptr, &out);
+    ASSERT_EQ(out.obj(), BSON("" << 1));
+}
+
 TEST(CollationIndexKeyTest, CollationAwareAppendReversesStringWithReverseMockCollator) {
     CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
     BSONObj dataObj = BSON("foo"
@@ -99,4 +114,44 @@ TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlySerializesWithEmbeddedN
     ASSERT_EQ(out.obj(), expectedObj);
 }
 
+TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesSimpleEmbeddedObject) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("" << BSON("a"
+                                      << "!foo"));
+    BSONObj expected = BSON("" << BSON("a"
+                                       << "oof!"));
+
+    BSONObjBuilder out;
+    CollationIndexKey::collationAwareIndexKeyAppend(dataObj.firstElement(), &collator, &out);
+    ASSERT_EQ(out.obj(), expected);
+}
+
+TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesSimpleEmbeddedArray) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = BSON("" << BSON_ARRAY("foo"
+                                            << "bar"));
+    BSONObj expected = BSON("" << BSON_ARRAY("oof"
+                                             << "rab"));
+
+    BSONObjBuilder out;
+    CollationIndexKey::collationAwareIndexKeyAppend(dataObj.firstElement(), &collator, &out);
+    ASSERT_EQ(out.obj(), expected);
+}
+
+TEST(CollationIndexKeyTest, CollationAwareAppendCorrectlyReversesComplexNesting) {
+    CollatorInterfaceMock collator(CollatorInterfaceMock::MockType::kReverseString);
+    BSONObj dataObj = fromjson(
+        "{ '' : [{'a': 'ha', 'b': 2},"
+        "'bar',"
+        "{'c': 2, 'd': 'ah', 'e': 'abc', 'f': ['cba', 'xyz']}]})");
+    BSONObj expected = fromjson(
+        "{ '' : [{'a': 'ah', 'b': 2},"
+        "'rab',"
+        "{'c': 2, 'd': 'ha', 'e': 'cba', 'f': ['abc', 'zyx']}]})");
+
+    BSONObjBuilder out;
+    CollationIndexKey::collationAwareIndexKeyAppend(dataObj.firstElement(), &collator, &out);
+    ASSERT_EQ(out.obj(), expected);
+}
+
 }  // namespace