SERVER-41730 Make TwoDAccessMethod generate keys using KeyString

6 files changed, 61 insertions, 14 deletions

diff --git a/src/mongo/db/geo/SConscript b/src/mongo/db/geo/SConscript
index 0feb1e4888a..fb13309cbfb 100644
--- a/src/mongo/db/geo/SConscript
+++ b/src/mongo/db/geo/SConscript
@@ -11,6 +11,7 @@ env.Library("geometry", [ "hash.cpp",
                           "r2_region_coverer.cpp" ],
             LIBDEPS = [ "$BUILD_DIR/mongo/base",
                         "$BUILD_DIR/mongo/db/common",
+                        '$BUILD_DIR/mongo/db/storage/key_string',
                         "$BUILD_DIR/third_party/s2/s2" ])
 
 # Geometry / BSON parsing and wrapping
diff --git a/src/mongo/db/geo/hash.cpp b/src/mongo/db/geo/hash.cpp
index f8ebed9a7da..6d870ece872 100644
--- a/src/mongo/db/geo/hash.cpp
+++ b/src/mongo/db/geo/hash.cpp
@@ -498,11 +498,31 @@ static void appendHashToBuilder(long long hash, BSONObjBuilder* builder, const c
     builder->appendBinData(fieldName, 8, bdtCustom, buf);
 }
 
+static void appendHashToKeyString(long long hash, KeyString::Builder* ks) {
+    char buf[8];
+#if MONGO_CONFIG_BYTE_ORDER == MONGO_LITTLE_ENDIAN
+    // Reverse the order of bytes when copying between BinData and GeoHash.
+    // GeoHashes are meant to be compared from MSB to LSB, where the first 2 MSB indicate the
+    // quadrant.
+    // In BinData, the GeoHash of a 2D index is compared from LSB to MSB, so the bytes should be
+    // reversed on little-endian systems
+    copyAndReverse(buf, (char*)&hash);
+#else
+    std::memcpy(buf, reinterpret_cast<char*>(&hash), 8);
+#endif
+    ks->appendBinData(BSONBinData(buf, 8, bdtCustom));
+}
+
 void GeoHash::appendHashMin(BSONObjBuilder* builder, const char* fieldName) const {
     // The min bound of a GeoHash region has all the unused suffix bits set to 0
     appendHashToBuilder(_hash, builder, fieldName);
 }
 
+void GeoHash::appendHashMin(KeyString::Builder* ks) const {
+    // The min bound of a GeoHash region has all the unused suffix bits set to 0
+    appendHashToKeyString(_hash, ks);
+}
+
 void GeoHash::appendHashMax(BSONObjBuilder* builder, const char* fieldName) const {
     // The max bound of a GeoHash region has all the unused suffix bits set to 1
     long long suffixMax = ~(geoBitSets.allX[_bits] | geoBitSets.allY[_bits]);
diff --git a/src/mongo/db/geo/hash.h b/src/mongo/db/geo/hash.h
index 0133f500e38..bfdca2cdffd 100644
--- a/src/mongo/db/geo/hash.h
+++ b/src/mongo/db/geo/hash.h
@@ -30,6 +30,7 @@
 #pragma once
 
 #include "mongo/db/jsobj.h"
+#include "mongo/db/storage/key_string.h"
 #include "mongo/platform/basic.h"
 
 namespace mongo {
@@ -111,6 +112,8 @@ public:
 
     // Append the minimum range of the hash to the builder provided (inclusive)
     void appendHashMin(BSONObjBuilder* builder, const char* fieldName) const;
+    // Append the minimum range of the hash to the KeyString provided (inclusive)
+    void appendHashMin(KeyString::Builder* ks) const;
     // Append the maximum range of the hash to the builder provided (inclusive)
     void appendHashMax(BSONObjBuilder* builder, const char* fieldName) const;
 
diff --git a/src/mongo/db/index/expression_keys_private.cpp b/src/mongo/db/index/expression_keys_private.cpp
index 847da8f61cb..ea3b3da6bac 100644
--- a/src/mongo/db/index/expression_keys_private.cpp
+++ b/src/mongo/db/index/expression_keys_private.cpp
@@ -270,8 +270,6 @@ void ExpressionKeysPrivate::get2DKeys(const BSONObj& obj,
                     continue;
             }
 
-            BSONObjBuilder b(64);
-
             // Stop if we don't need to get anything but location objects
             if (!keys) {
                 if (singleElement)
@@ -280,7 +278,8 @@ void ExpressionKeysPrivate::get2DKeys(const BSONObj& obj,
                     continue;
             }
 
-            params.geoHashConverter->hash(locObj, &obj).appendHashMin(&b, "");
+            KeyString::Builder keyString(keyStringVersion, ordering);
+            params.geoHashConverter->hash(locObj, &obj).appendHashMin(&keyString);
 
             // Go through all the other index keys
             for (vector<pair<string, int>>::const_iterator i = params.other.begin();
@@ -291,21 +290,15 @@ void ExpressionKeysPrivate::get2DKeys(const BSONObj& obj,
                 dps::extractAllElementsAlongPath(obj, i->first, eSet);
 
                 if (eSet.size() == 0)
-                    b.appendNull("");
+                    keyString.appendNull();
                 else if (eSet.size() == 1)
-                    b.appendAs(*(eSet.begin()), "");
+                    keyString.appendBSONElement(*(eSet.begin()));
                 else {
                     // If we have more than one key, store as an array of the objects
-                    BSONArrayBuilder aBuilder;
-
-                    for (BSONElementSet::iterator ei = eSet.begin(); ei != eSet.end(); ++ei) {
-                        aBuilder.append(*ei);
-                    }
-
-                    b.append("", aBuilder.arr());
+                    keyString.appendSetAsArray(eSet);
                 }
             }
-            KeyString::Builder keyString(keyStringVersion, b.obj(), ordering);
+
             if (id) {
                 keyString.appendRecordId(*id);
             }
diff --git a/src/mongo/db/storage/key_string.cpp b/src/mongo/db/storage/key_string.cpp
index 5f018544e96..58daae6b73c 100644
--- a/src/mongo/db/storage/key_string.cpp
+++ b/src/mongo/db/storage/key_string.cpp
@@ -384,6 +384,20 @@ void BuilderBase<BufferT>::appendUndefined() {
 }
 
 template <class BufferT>
+void BuilderBase<BufferT>::appendBinData(const BSONBinData& data) {
+    _verifyAppendingState();
+    _appendBinData(data, _shouldInvertOnAppend());
+    _elemCount++;
+}
+
+template <class BufferT>
+void BuilderBase<BufferT>::appendSetAsArray(const BSONElementSet& set, const StringTransformFn& f) {
+    _verifyAppendingState();
+    _appendSetAsArray(set, _shouldInvertOnAppend(), nullptr);
+    _elemCount++;
+}
+
+template <class BufferT>
 void BuilderBase<BufferT>::_appendDiscriminator(const Discriminator discriminator) {
     // The discriminator forces this KeyString to compare Less/Greater than any KeyString with
     // the same prefix of keys. As an example, this can be used to land on the first key in the
@@ -594,7 +608,19 @@ void BuilderBase<BufferT>::_appendArray(const BSONArray& val,
                                         bool invert,
                                         const StringTransformFn& f) {
     _append(CType::kArray, invert);
-    BSONForEach(elem, val) {
+    for (const auto& elem : val) {
+        // No generic ctype byte needed here since no name is encoded.
+        _appendBsonValue(elem, invert, nullptr, f);
+    }
+    _append(int8_t(0), invert);
+}
+
+template <class BufferT>
+void BuilderBase<BufferT>::_appendSetAsArray(const BSONElementSet& val,
+                                             bool invert,
+                                             const StringTransformFn& f) {
+    _append(CType::kArray, invert);
+    for (const auto& elem : val) {
         // No generic ctype byte needed here since no name is encoded.
         _appendBsonValue(elem, invert, nullptr, f);
     }
diff --git a/src/mongo/db/storage/key_string.h b/src/mongo/db/storage/key_string.h
index da137362c16..ae58d48261c 100644
--- a/src/mongo/db/storage/key_string.h
+++ b/src/mongo/db/storage/key_string.h
@@ -32,6 +32,7 @@
 #include <limits>
 
 #include "mongo/base/static_assert.h"
+#include "mongo/bson/bsonelement_comparator_interface.h"
 #include "mongo/bson/bsonmisc.h"
 #include "mongo/bson/bsonobj.h"
 #include "mongo/bson/bsonobjbuilder.h"
@@ -455,6 +456,8 @@ public:
     void appendNumberLong(long long num);
     void appendNull();
     void appendUndefined();
+    void appendBinData(const BSONBinData& data);
+    void appendSetAsArray(const BSONElementSet& set, const StringTransformFn& f = nullptr);
 
     /**
      * Resets to an empty state.
@@ -537,6 +540,7 @@ private:
     void _appendRegex(const BSONRegEx& val, bool invert);
     void _appendDBRef(const BSONDBRef& val, bool invert);
     void _appendArray(const BSONArray& val, bool invert, const StringTransformFn& f);
+    void _appendSetAsArray(const BSONElementSet& val, bool invert, const StringTransformFn& f);
     void _appendObject(const BSONObj& val, bool invert, const StringTransformFn& f);
     void _appendNumberDouble(const double num, bool invert);
     void _appendNumberLong(const long long num, bool invert);