SERVER-41721 Make IndexAccessMethod::getKeys output a KeyString set

3 files changed, 100 insertions, 50 deletions

diff --git a/src/mongo/db/fts/fts_index_format.cpp b/src/mongo/db/fts/fts_index_format.cpp
index ab9950635fb..b514a67a6dd 100644
--- a/src/mongo/db/fts/fts_index_format.cpp
+++ b/src/mongo/db/fts/fts_index_format.cpp
@@ -135,7 +135,12 @@ MONGO_INITIALIZER(FTSIndexFormat)(InitializerContext* context) {
     return Status::OK();
 }
 
-void FTSIndexFormat::getKeys(const FTSSpec& spec, const BSONObj& obj, BSONObjSet* keys) {
+void FTSIndexFormat::getKeys(const FTSSpec& spec,
+                             const BSONObj& obj,
+                             KeyStringSet* keys,
+                             KeyString::Version keyStringVersion,
+                             Ordering ordering,
+                             boost::optional<RecordId> id) {
     int extraSize = 0;
     vector<BSONElement> extrasBefore;
     vector<BSONElement> extrasAfter;
@@ -182,7 +187,11 @@ void FTSIndexFormat::getKeys(const FTSSpec& spec, const BSONObj& obj, BSONObjSet
 
         verify(guess >= res.objsize());
 
-        keys->insert(res);
+        KeyString::HeapBuilder keyString(keyStringVersion, res, ordering);
+        if (id) {
+            keyString.appendRecordId(*id);
+        }
+        keys->insert(keyString.release());
         keyBSONSize += res.objsize();
     }
 }
diff --git a/src/mongo/db/fts/fts_index_format.h b/src/mongo/db/fts/fts_index_format.h
index dd83e8603a8..5ac0823c0c0 100644
--- a/src/mongo/db/fts/fts_index_format.h
+++ b/src/mongo/db/fts/fts_index_format.h
@@ -34,6 +34,8 @@
 #include "mongo/base/string_data.h"
 #include "mongo/bson/bsonobj_comparator_interface.h"
 #include "mongo/db/fts/fts_util.h"
+#include "mongo/db/storage/key_string.h"
+#include "mongo/db/storage/sorted_data_interface.h"
 
 namespace mongo {
 
@@ -43,7 +45,12 @@ class FTSSpec;
 
 class FTSIndexFormat {
 public:
-    static void getKeys(const FTSSpec& spec, const BSONObj& document, BSONObjSet* keys);
+    static void getKeys(const FTSSpec& spec,
+                        const BSONObj& document,
+                        KeyStringSet* keys,
+                        KeyString::Version keyStringVersion,
+                        Ordering ordering,
+                        boost::optional<RecordId> id = boost::none);
 
     /**
      * Helper method to get return entry from the FTSIndex as a BSONObj
diff --git a/src/mongo/db/fts/fts_index_format_test.cpp b/src/mongo/db/fts/fts_index_format_test.cpp
index c9d6779e639..46b10c81228 100644
--- a/src/mongo/db/fts/fts_index_format_test.cpp
+++ b/src/mongo/db/fts/fts_index_format_test.cpp
@@ -51,15 +51,17 @@ using unittest::assertGet;
 TEST(FTSIndexFormat, Simple1) {
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << BSON("data"
                                                                << "text")))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     FTSIndexFormat::getKeys(spec,
                             BSON("data"
                                  << "cat sat"),
-                            &keys);
+                            &keys,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
 
     ASSERT_EQUALS(2U, keys.size());
-    for (BSONObjSet::const_iterator i = keys.begin(); i != keys.end(); ++i) {
-        BSONObj key = *i;
+    for (auto& keyString : keys) {
+        auto key = KeyString::toBson(keyString, Ordering::make(BSONObj()));
         ASSERT_EQUALS(2, key.nFields());
         ASSERT_EQUALS(String, key.firstElement().type());
     }
@@ -69,15 +71,17 @@ TEST(FTSIndexFormat, ExtraBack1) {
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << BSON("data"
                                                                << "text"
                                                                << "x" << 1)))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     FTSIndexFormat::getKeys(spec,
                             BSON("data"
                                  << "cat"
                                  << "x" << 5),
-                            &keys);
+                            &keys,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
 
     ASSERT_EQUALS(1U, keys.size());
-    BSONObj key = *(keys.begin());
+    auto key = KeyString::toBson(*keys.begin(), Ordering::make(BSONObj()));
     ASSERT_EQUALS(3, key.nFields());
     BSONObjIterator i(key);
     ASSERT_EQUALS(StringData("cat"), i.next().valuestr());
@@ -88,15 +92,17 @@ TEST(FTSIndexFormat, ExtraBack1) {
 TEST(FTSIndexFormat, ExtraFront1) {
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << BSON("x" << 1 << "data"
                                                                    << "text")))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     FTSIndexFormat::getKeys(spec,
                             BSON("data"
                                  << "cat"
                                  << "x" << 5),
-                            &keys);
+                            &keys,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
 
     ASSERT_EQUALS(1U, keys.size());
-    BSONObj key = *(keys.begin());
+    auto key = KeyString::toBson(*keys.begin(), Ordering::make(BSONObj()));
     ASSERT_EQUALS(3, key.nFields());
     BSONObjIterator i(key);
     ASSERT_EQUALS(5, i.next().numberInt());
@@ -108,18 +114,22 @@ TEST(FTSIndexFormat, StopWords1) {
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << BSON("data"
                                                                << "text")))));
 
-    BSONObjSet keys1 = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys1;
     FTSIndexFormat::getKeys(spec,
                             BSON("data"
                                  << "computer"),
-                            &keys1);
+                            &keys1,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
     ASSERT_EQUALS(1U, keys1.size());
 
-    BSONObjSet keys2 = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys2;
     FTSIndexFormat::getKeys(spec,
                             BSON("data"
                                  << "any computer"),
-                            &keys2);
+                            &keys2,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
     ASSERT_EQUALS(1U, keys2.size());
 }
 
@@ -127,10 +137,10 @@ TEST(FTSIndexFormat, StopWords1) {
  * Helper function to compare keys returned in getKeys() result
  * with expected values.
  */
-void assertEqualsIndexKeys(std::set<std::string>& expectedKeys, const BSONObjSet& keys) {
+void assertEqualsIndexKeys(std::set<std::string>& expectedKeys, const KeyStringSet& keys) {
     ASSERT_EQUALS(expectedKeys.size(), keys.size());
-    for (BSONObjSet::const_iterator i = keys.begin(); i != keys.end(); ++i) {
-        BSONObj key = *i;
+    for (auto& keyString : keys) {
+        auto key = KeyString::toBson(keyString, Ordering::make(BSONObj()));
         ASSERT_EQUALS(2, key.nFields());
         ASSERT_EQUALS(String, key.firstElement().type());
         string s = key.firstElement().String();
@@ -156,14 +166,18 @@ TEST(FTSIndexFormat, LongWordsTextIndexVersion1) {
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << BSON("data"
                                                                << "text")
                                                        << "textIndexVersion" << 1))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     string longPrefix(1024U, 'a');
     // "aaa...aaacat"
     string longWordCat = longPrefix + "cat";
     // "aaa...aaasat"
     string longWordSat = longPrefix + "sat";
     string text = str::stream() << longWordCat << " " << longWordSat;
-    FTSIndexFormat::getKeys(spec, BSON("data" << text), &keys);
+    FTSIndexFormat::getKeys(spec,
+                            BSON("data" << text),
+                            &keys,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
 
     // Hard-coded expected computed keys for future-proofing.
     std::set<string> expectedKeys;
@@ -185,7 +199,7 @@ TEST(FTSIndexFormat, LongWordTextIndexVersion2) {
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << BSON("data"
                                                                << "text")
                                                        << "textIndexVersion" << 2))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     string longPrefix(1024U, 'a');
     // "aaa...aaacat"
     string longWordCat = longPrefix + "cat";
@@ -194,7 +208,11 @@ TEST(FTSIndexFormat, LongWordTextIndexVersion2) {
     // "aaa...aaamongodbfts"
     string longWordMongoDBFts = longPrefix + "mongodbfts";
     string text = str::stream() << longWordCat << " " << longWordSat << " " << longWordMongoDBFts;
-    FTSIndexFormat::getKeys(spec, BSON("data" << text), &keys);
+    FTSIndexFormat::getKeys(spec,
+                            BSON("data" << text),
+                            &keys,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
 
     // Hard-coded expected computed keys for future-proofing.
     std::set<string> expectedKeys;
@@ -218,14 +236,18 @@ TEST(FTSIndexFormat, LongWordTextIndexVersion3) {
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << BSON("data"
                                                                << "text")
                                                        << "textIndexVersion" << 3))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     string longPrefix(1024U, 'a');
     // "aaa...aaacat"
     string longWordCat = longPrefix + "cat";
     // "aaa...aaasat"
     string longWordSat = longPrefix + "sat";
     string text = str::stream() << longWordCat << " " << longWordSat;
-    FTSIndexFormat::getKeys(spec, BSON("data" << text), &keys);
+    FTSIndexFormat::getKeys(spec,
+                            BSON("data" << text),
+                            &keys,
+                            KeyString::Version::kLatestVersion,
+                            Ordering::make(BSONObj()));
 
     // Hard-coded expected computed keys for future-proofing.
     std::set<string> expectedKeys;
@@ -246,63 +268,74 @@ TEST(FTSIndexFormat, LongWordTextIndexVersion3) {
 TEST(FTSIndexFormat, GetKeysWithLeadingEmptyArrayThrows) {
     BSONObj keyPattern = fromjson("{'a.b': 1, data: 'text'}");
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << keyPattern << "textIndexVersion" << 3))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     BSONObj objToIndex = fromjson("{a: {b: []}, data: 'foo'}");
-    ASSERT_THROWS_CODE(FTSIndexFormat::getKeys(spec, objToIndex, &keys),
-                       AssertionException,
-                       ErrorCodes::CannotBuildIndexKeys);
+    ASSERT_THROWS_CODE(
+        FTSIndexFormat::getKeys(
+            spec, objToIndex, &keys, KeyString::Version::kLatestVersion, Ordering::make(BSONObj())),
+        AssertionException,
+        ErrorCodes::CannotBuildIndexKeys);
 }
 
 TEST(FTSIndexFormat, GetKeysWithTrailingEmptyArrayThrows) {
     BSONObj keyPattern = fromjson("{data: 'text', 'a.b': 1}");
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << keyPattern << "textIndexVersion" << 3))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     BSONObj objToIndex = fromjson("{a: {b: []}, data: 'foo'}");
-    ASSERT_THROWS_CODE(FTSIndexFormat::getKeys(spec, objToIndex, &keys),
-                       AssertionException,
-                       ErrorCodes::CannotBuildIndexKeys);
+    ASSERT_THROWS_CODE(
+        FTSIndexFormat::getKeys(
+            spec, objToIndex, &keys, KeyString::Version::kLatestVersion, Ordering::make(BSONObj())),
+        AssertionException,
+        ErrorCodes::CannotBuildIndexKeys);
 }
 
 TEST(FTSIndexFormat, GetKeysWithLeadingSingleElementArrayThrows) {
     BSONObj keyPattern = fromjson("{'a.b': 1, data: 'text'}");
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << keyPattern << "textIndexVersion" << 3))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     BSONObj objToIndex = fromjson("{a: [{b: 9}], data: 'foo'}");
-    ASSERT_THROWS_CODE(FTSIndexFormat::getKeys(spec, objToIndex, &keys),
-                       AssertionException,
-                       ErrorCodes::CannotBuildIndexKeys);
+    ASSERT_THROWS_CODE(
+        FTSIndexFormat::getKeys(
+            spec, objToIndex, &keys, KeyString::Version::kLatestVersion, Ordering::make(BSONObj())),
+        AssertionException,
+        ErrorCodes::CannotBuildIndexKeys);
 }
 
 TEST(FTSIndexFormat, GetKeysWithTrailingSingleElementArrayThrows) {
     BSONObj keyPattern = fromjson("{data: 'text', 'a.b': 1}");
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << keyPattern << "textIndexVersion" << 3))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     BSONObj objToIndex = fromjson("{a: [{b: 9}], data: 'foo'}");
-    ASSERT_THROWS_CODE(FTSIndexFormat::getKeys(spec, objToIndex, &keys),
-                       AssertionException,
-                       ErrorCodes::CannotBuildIndexKeys);
+    ASSERT_THROWS_CODE(
+        FTSIndexFormat::getKeys(
+            spec, objToIndex, &keys, KeyString::Version::kLatestVersion, Ordering::make(BSONObj())),
+        AssertionException,
+        ErrorCodes::CannotBuildIndexKeys);
 }
 
 TEST(FTSIndexFormat, GetKeysWithMultiElementArrayThrows) {
     BSONObj keyPattern = fromjson("{'a.b': 1, 'a.c': 'text'}");
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << keyPattern << "textIndexVersion" << 3))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     BSONObj objToIndex = fromjson("{a: [{b: 9, c: 'foo'}, {b: 10, c: 'bar'}]}");
-    ASSERT_THROWS_CODE(FTSIndexFormat::getKeys(spec, objToIndex, &keys),
-                       AssertionException,
-                       ErrorCodes::CannotBuildIndexKeys);
+    ASSERT_THROWS_CODE(
+        FTSIndexFormat::getKeys(
+            spec, objToIndex, &keys, KeyString::Version::kLatestVersion, Ordering::make(BSONObj())),
+        AssertionException,
+        ErrorCodes::CannotBuildIndexKeys);
 }
 
 TEST(FTSIndexFormat, GetKeysWithPositionalPathAllowed) {
     BSONObj keyPattern = fromjson("{'a.0': 1, 'a.b': 'text'}");
     FTSSpec spec(assertGet(FTSSpec::fixSpec(BSON("key" << keyPattern << "textIndexVersion" << 3))));
-    BSONObjSet keys = SimpleBSONObjComparator::kInstance.makeBSONObjSet();
+    KeyStringSet keys;
     BSONObj objToIndex = fromjson("{a: [{b: 'foo'}, {b: 'bar'}]}");
-    FTSIndexFormat::getKeys(spec, objToIndex, &keys);
+    FTSIndexFormat::getKeys(
+        spec, objToIndex, &keys, KeyString::Version::kLatestVersion, Ordering::make(BSONObj()));
     ASSERT_EQ(2U, keys.size());
 
     {
-        BSONObj key = *(keys.begin());
+        auto key = KeyString::toBson(*keys.begin(), Ordering::make(BSONObj()));
         ASSERT_EQ(3, key.nFields());
         BSONObjIterator it{key};
         ASSERT_BSONELT_EQ(it.next(), fromjson("{'': {b: 'foo'}}").firstElement());
@@ -310,7 +343,8 @@ TEST(FTSIndexFormat, GetKeysWithPositionalPathAllowed) {
     }
 
     {
-        BSONObj key = *(++keys.begin());
+        auto next = ++keys.begin();
+        auto key = KeyString::toBson(*next, Ordering::make(BSONObj()));
         ASSERT_EQ(3, key.nFields());
         BSONObjIterator it{key};
         ASSERT_BSONELT_EQ(it.next(), fromjson("{'': {b: 'foo'}}").firstElement());