diff options
author | David Hatch <david.hatch@mongodb.com> | 2016-06-14 11:51:33 -0400 |
---|---|---|
committer | David Hatch <david.hatch@mongodb.com> | 2016-06-22 15:42:59 -0400 |
commit | 40f20eca105a5e06a72df583ac654f946e9b058e (patch) | |
tree | 0ca485f5a24297236b9455a6698e7135802e08cb /src/mongo/db/query/collation | |
parent | f517b05141a3f554d2fe51838ed33bc98cb8c5f2 (diff) | |
download | mongo-40f20eca105a5e06a72df583ac654f946e9b058e.tar.gz |
SERVER-23172 Allow use of indices for collation-aware queries that match nested objects or arrays.
Diffstat (limited to 'src/mongo/db/query/collation')
-rw-r--r-- | src/mongo/db/query/collation/collation_index_key.cpp | 81 | ||||
-rw-r--r-- | src/mongo/db/query/collation/collation_index_key.h | 28 | ||||
-rw-r--r-- | src/mongo/db/query/collation/collation_index_key_test.cpp | 63 |
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 |