Merge pull request #15291 from vespa-engine/toregge/add-gid-to-lid-map-key-class-for-document-meta-store

Add 32 most significant bits of gid key to key in btree mapping from

12 files changed, 137 insertions, 60 deletions

diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/CMakeLists.txt b/searchcore/src/vespa/searchcore/proton/documentmetastore/CMakeLists.txt
index 10e0c149ad2..be298f87cf1 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/CMakeLists.txt
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/CMakeLists.txt
@@ -9,6 +9,7 @@ vespa_add_library(searchcore_documentmetastore STATIC
     documentmetastoreflushtarget.cpp
     documentmetastoreinitializer.cpp
     documentmetastoresaver.cpp
+    gid_to_lid_map_key.cpp
     search_context.cpp
     lid_allocator.cpp
     lid_gid_key_comparator.cpp
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.cpp b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.cpp
index b27a5d57fc3..b3dfebfa9c0 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.cpp
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.cpp
@@ -28,6 +28,7 @@ LOG_SETUP(".proton.documentmetastore");
 using document::BucketId;
 using document::GlobalId;
 using proton::bucketdb::BucketState;
+using proton::documentmetastore::GidToLidMapKey;
 using search::AttributeVector;
 using search::FileReader;
 using search::GrowStrategy;
@@ -174,11 +175,12 @@ DocumentMetaStore::ensureSpace(DocId lid)
 }
 
 void
-DocumentMetaStore::insert(DocId lid, const RawDocumentMetaData &metaData)
+DocumentMetaStore::insert(GidToLidMapKey key, const RawDocumentMetaData &metaData)
 {
+    auto lid = key.get_lid();
     ensureSpace(lid);
     _metaDataStore[lid] = metaData;
-    _gidToLidMap.insert(_gid_to_lid_map_write_itr, lid, BTreeNoLeafData());
+    _gidToLidMap.insert(_gid_to_lid_map_write_itr, key, BTreeNoLeafData());
     // flush writes to meta store rcu vector before new entry is visible
     // from frozen root or lid based scan
     std::atomic_thread_fence(std::memory_order_release);
@@ -249,7 +251,7 @@ DocumentMetaStore::readNextDoc(documentmetastore::Reader & reader, TreeType::Bui
     meta.setBucketUsedBits(reader.getNextBucketUsedBits());
     meta.setDocSize(reader.getNextDocSize());
     meta.setTimestamp(reader.getNextTimestamp());
-    treeBuilder.insert(lid, BTreeNoLeafData());
+    treeBuilder.insert(GidToLidMapKey(lid, meta.getGid()), BTreeNoLeafData());
     assert(!validLid(lid));
     _lidAlloc.registerLid(lid);
     return lid;
@@ -305,7 +307,8 @@ DocumentMetaStore::lowerBound(const BucketId &bucketId,
 {
     document::GlobalId first(document::GlobalId::calculateFirstInBucket(bucketId));
     KeyComp lowerComp(first, _metaDataStore);
-    return treeView.lowerBound(KeyComp::FIND_DOC_ID, lowerComp);
+    auto find_key = GidToLidMapKey::make_find_key(first);
+    return treeView.lowerBound(find_key, lowerComp);
 }
 
 template <typename TreeView>
@@ -315,7 +318,8 @@ DocumentMetaStore::upperBound(const BucketId &bucketId,
 {
     document::GlobalId last(document::GlobalId::calculateLastInBucket(bucketId));
     KeyComp upperComp(last, _metaDataStore);
-    return treeView.upperBound(KeyComp::FIND_DOC_ID, upperComp);
+    auto find_key = GidToLidMapKey::make_find_key(last);
+    return treeView.upperBound(find_key, upperComp);
 }
 
 void
@@ -361,7 +365,7 @@ DocumentMetaStore::unload()
     BucketId prev;
     BucketState prevDelta;
     for (; itr.valid(); ++itr) {
-        uint32_t lid = itr.getKey();
+        uint32_t lid = itr.getKey().get_lid();
         assert(validLid(lid));
         RawDocumentMetaData &metaData = _metaDataStore[lid];
         BucketId bucketId = metaData.getBucketId();
@@ -421,12 +425,13 @@ DocumentMetaStore::inspectExisting(const GlobalId &gid, uint64_t prepare_serial_
     assert(_lidAlloc.isFreeListConstructed());
     Result res;
     KeyComp comp(gid, _metaDataStore);
+    auto find_key = GidToLidMapKey::make_find_key(gid);
     auto& itr = _gid_to_lid_map_write_itr;
-    itr.lower_bound(_gidToLidMap.getRoot(), KeyComp::FIND_DOC_ID, comp);
+    itr.lower_bound(_gidToLidMap.getRoot(), find_key, comp);
     _gid_to_lid_map_write_itr_prepare_serial_num = prepare_serial_num;
-    bool found = itr.valid() && !comp(KeyComp::FIND_DOC_ID, itr.getKey());
+    bool found = itr.valid() && !comp(find_key, itr.getKey());
     if (found) {
-        res.setLid(itr.getKey());
+        res.setLid(itr.getKey().get_lid());
         res.fillPrev(_metaDataStore[res.getLid()].getTimestamp());
         res.markSuccess();
     }
@@ -439,16 +444,17 @@ DocumentMetaStore::inspect(const GlobalId &gid, uint64_t prepare_serial_num)
     assert(_lidAlloc.isFreeListConstructed());
     Result res;
     KeyComp comp(gid, _metaDataStore);
+    auto find_key = GidToLidMapKey::make_find_key(gid);
     auto& itr = _gid_to_lid_map_write_itr;
-    itr.lower_bound(_gidToLidMap.getRoot(), KeyComp::FIND_DOC_ID, comp);
+    itr.lower_bound(_gidToLidMap.getRoot(), find_key, comp);
     _gid_to_lid_map_write_itr_prepare_serial_num = prepare_serial_num;
-    bool found = itr.valid() && !comp(KeyComp::FIND_DOC_ID, itr.getKey());
+    bool found = itr.valid() && !comp(find_key, itr.getKey());
     if (!found) {
         DocId myLid = peekFreeLid();
         res.setLid(myLid);
         res.markSuccess();
     } else {
-        res.setLid(itr.getKey());
+        res.setLid(itr.getKey().get_lid());
         res.fillPrev(_metaDataStore[res.getLid()].getTimestamp());
         res.markSuccess();
     }
@@ -466,11 +472,12 @@ DocumentMetaStore::put(const GlobalId &gid,
     Result res;
     RawDocumentMetaData metaData(gid, bucketId, timestamp, docSize);
     KeyComp comp(metaData, _metaDataStore);
+    auto find_key = GidToLidMapKey::make_find_key(gid);
     auto& itr = _gid_to_lid_map_write_itr;
     if (prepare_serial_num == 0u || _gid_to_lid_map_write_itr_prepare_serial_num != prepare_serial_num) {
-        itr.lower_bound(_gidToLidMap.getRoot(), KeyComp::FIND_DOC_ID, comp);
+        itr.lower_bound(_gidToLidMap.getRoot(), find_key, comp);
     }
-    bool found = itr.valid() && !comp(KeyComp::FIND_DOC_ID, itr.getKey());
+    bool found = itr.valid() && !comp(find_key, itr.getKey());
     if (!found) {
         if (validLid(lid)) {
             throw IllegalStateException(
@@ -490,10 +497,10 @@ DocumentMetaStore::put(const GlobalId &gid,
             assert(freeLid == lid);
             (void) freeLid;
         }
-        insert(lid, metaData);
+        insert(GidToLidMapKey(lid, find_key.get_gid_key()), metaData);
         res.setLid(lid);
         res.markSuccess();
-    } else if (lid != itr.getKey()) {
+    } else if (lid != itr.getKey().get_lid()) {
         throw IllegalStateException(
                 make_string(
                         "document meta data store"
@@ -503,7 +510,7 @@ DocumentMetaStore::put(const GlobalId &gid,
                         " gid found, but using another lid '%u'",
                         lid,
                         gid.toString().c_str(),
-                        itr.getKey()));
+                        itr.getKey().get_lid()));
     } else {
         res.setLid(lid);
         res.fillPrev(_metaDataStore[lid].getTimestamp());
@@ -541,11 +548,12 @@ DocumentMetaStore::remove(DocId lid, uint64_t prepare_serial_num, BucketDBOwner:
 {
     const GlobalId & gid = getRawGid(lid);
     KeyComp comp(gid, _metaDataStore);
+    GidToLidMapKey find_key(lid, gid);
     auto& itr = _gid_to_lid_map_write_itr;
     if (prepare_serial_num == 0u || _gid_to_lid_map_write_itr_prepare_serial_num != prepare_serial_num) {
-        itr.lower_bound(_gidToLidMap.getRoot(), lid, comp);
+        itr.lower_bound(_gidToLidMap.getRoot(), find_key, comp);
     }
-    if (!itr.valid() || comp(lid, itr.getKey())) {
+    if (!itr.valid() || comp(find_key, itr.getKey())) {
         throw IllegalStateException(make_string(
                         "document meta data store corrupted,"
                         " cannot remove"
@@ -598,14 +606,15 @@ DocumentMetaStore::move(DocId fromLid, DocId toLid, uint64_t prepare_serial_num)
     _metaDataStore[toLid] = _metaDataStore[fromLid];
     const GlobalId & gid = getRawGid(fromLid);
     KeyComp comp(gid, _metaDataStore);
+    GidToLidMapKey find_key(fromLid, gid);
     auto& itr = _gid_to_lid_map_write_itr;
     if (prepare_serial_num == 0u || _gid_to_lid_map_write_itr_prepare_serial_num != prepare_serial_num) {
-        itr.lower_bound(_gidToLidMap.getRoot(), fromLid, comp);
+        itr.lower_bound(_gidToLidMap.getRoot(), find_key, comp);
     }
     assert(itr.valid());
-    assert(itr.getKey() == fromLid);
+    assert(itr.getKey().get_lid() == fromLid);
     _gidToLidMap.thaw(itr);
-    itr.writeKey(toLid);
+    itr.writeKey(GidToLidMapKey(toLid, find_key.get_gid_key()));
     _lidAlloc.moveLidEnd(fromLid, toLid);
     incGeneration();
 }
@@ -665,12 +674,13 @@ DocumentMetaStore::getLid(const GlobalId &gid, DocId &lid) const
 {
     GlobalId value(gid);
     KeyComp comp(value, _metaDataStore);
+    auto find_key = GidToLidMapKey::make_find_key(gid);
     TreeType::ConstIterator itr =
-        _gidToLidMap.getFrozenView().find(KeyComp::FIND_DOC_ID, comp);
+        _gidToLidMap.getFrozenView().find(find_key, comp);
     if (!itr.valid()) {
         return false;
     }
-    lid = itr.getKey();
+    lid = itr.getKey().get_lid();
     return true;
 }
 
@@ -707,7 +717,7 @@ DocumentMetaStore::getMetaData(const BucketId &bucketId,
     TreeType::ConstIterator itr = lowerBound(bucketId, frozenTreeView);
     TreeType::ConstIterator end = upperBound(bucketId, frozenTreeView);
     for (; itr != end; ++itr) {
-        DocId lid = itr.getKey();
+        DocId lid = itr.getKey().get_lid();
         if (validLid(lid)) {
             const RawDocumentMetaData &rawData = getRawMetaData(lid);
             if (bucketId.getUsedBits() != rawData.getBucketUsedBits())
@@ -779,7 +789,8 @@ DocumentMetaStore::lowerBound(const GlobalId &gid) const
 {
     // Called by writer thread
     KeyComp comp(gid, _metaDataStore);
-    return _gidToLidMap.lowerBound(KeyComp::FIND_DOC_ID, comp);
+    auto find_key = GidToLidMapKey::make_find_key(gid);
+    return _gidToLidMap.lowerBound(find_key, comp);
 }
 
 DocumentMetaStore::Iterator
@@ -787,7 +798,8 @@ DocumentMetaStore::upperBound(const GlobalId &gid) const
 {
     // Called by writer thread
     KeyComp comp(gid, _metaDataStore);
-    return _gidToLidMap.upperBound(KeyComp::FIND_DOC_ID, comp);
+    auto find_key = GidToLidMapKey::make_find_key(gid);
+    return _gidToLidMap.upperBound(find_key, comp);
 }
 
 void
@@ -797,7 +809,7 @@ DocumentMetaStore::getLids(const BucketId &bucketId, std::vector<DocId> &lids)
     TreeType::Iterator itr = lowerBound(bucketId);
     TreeType::Iterator end = upperBound(bucketId);
     for (; itr != end; ++itr) {
-        DocId lid = itr.getKey();
+        DocId lid = itr.getKey().get_lid();
         assert(validLid(lid));
         const RawDocumentMetaData &metaData = getRawMetaData(lid);
         uint8_t bucketUsedBits = metaData.getBucketUsedBits();
@@ -828,7 +840,7 @@ DocumentMetaStore::handleSplit(const bucketdb::SplitBucketSession &session)
     TreeType::Iterator end = upperBound(source);
     bucketdb::BucketDeltaPair deltas;
     for (; itr != end; ++itr) {
-        DocId lid = itr.getKey();
+        DocId lid = itr.getKey().get_lid();
         assert(validLid(lid));
         RawDocumentMetaData &metaData = _metaDataStore[lid];
         uint8_t bucketUsedBits = metaData.getBucketUsedBits();
@@ -873,7 +885,7 @@ DocumentMetaStore::handleJoin(const bucketdb::JoinBucketsSession &session)
     TreeType::Iterator end = upperBound(target);
     bucketdb::BucketDeltaPair deltas;
     for (; itr != end; ++itr) {
-        DocId lid = itr.getKey();
+        DocId lid = itr.getKey().get_lid();
         assert(validLid(lid));
         RawDocumentMetaData &metaData = _metaDataStore[lid];
         assert(BucketId::validUsedBits(metaData.getBucketUsedBits()));
@@ -913,7 +925,7 @@ DocumentMetaStore::updateActiveLids(const BucketId &bucketId, bool active)
     TreeType::Iterator end = upperBound(bucketId);
     uint8_t bucketUsedBits = bucketId.getUsedBits();
     for (; itr != end; ++itr) {
-        DocId lid = itr.getKey();
+        DocId lid = itr.getKey().get_lid();
         assert(validLid(lid));
         RawDocumentMetaData &metaData = _metaDataStore[lid];
         if (metaData.getBucketUsedBits() != bucketUsedBits) {
@@ -1038,18 +1050,18 @@ DocumentMetaStore::getVersion() const
 void
 DocumentMetaStore::foreach(const search::IGidToLidMapperVisitor &visitor) const
 {
-    beginFrozen().foreach_key([this,&visitor](uint32_t lid)
-                              { visitor.visit(getRawMetaData(lid).getGid(), lid); });
+    beginFrozen().foreach_key([this,&visitor](GidToLidMapKey key)
+                              { visitor.visit(getRawMetaData(key.get_lid()).getGid(), key.get_lid()); });
 }
 
 }  // namespace proton
 
 namespace vespalib::btree {
 
-template class BTreeIteratorBase<proton::DocumentMetaStore::DocId, BTreeNoLeafData, NoAggregated, BTreeDefaultTraits::INTERNAL_SLOTS, BTreeDefaultTraits::LEAF_SLOTS, BTreeDefaultTraits::PATH_SIZE>;
+template class BTreeIteratorBase<proton::documentmetastore::GidToLidMapKey, BTreeNoLeafData, NoAggregated, BTreeDefaultTraits::INTERNAL_SLOTS, BTreeDefaultTraits::LEAF_SLOTS, BTreeDefaultTraits::PATH_SIZE>;
 
-template class BTreeConstIterator<proton::DocumentMetaStore::DocId, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
+template class BTreeConstIterator<proton::documentmetastore::GidToLidMapKey, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
 
-template class BTreeIterator<proton::DocumentMetaStore::DocId, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
+template class BTreeIterator<proton::documentmetastore::GidToLidMapKey, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
 
 }
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.h b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.h
index ccb44a0a452..74447878f09 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.h
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastore.h
@@ -60,7 +60,7 @@ private:
     // Lids are stored as keys in the tree, sorted by their gid
     // counterpart.  The LidGidKeyComparator class maps from lids -> metadata by
     // using the metadata store.
-    typedef vespalib::btree::BTree<DocId, vespalib::btree::BTreeNoLeafData,
+    typedef vespalib::btree::BTree<documentmetastore::GidToLidMapKey, vespalib::btree::BTreeNoLeafData,
                                  vespalib::btree::NoAggregated, const KeyComp &> TreeType;
 
     MetaDataStore       _metaDataStore;
@@ -77,7 +77,7 @@ private:
     DocId getFreeLid();
     DocId peekFreeLid();
     VESPA_DLL_LOCAL void ensureSpace(DocId lid);
-    void insert(DocId lid, const RawDocumentMetaData &metaData);
+    void insert(documentmetastore::GidToLidMapKey key, const RawDocumentMetaData &metaData);
 
     const GlobalId & getRawGid(DocId lid) const { return getRawMetaData(lid).getGid(); }
 
@@ -265,10 +265,10 @@ public:
 
 namespace vespalib::btree {
 
-extern template class BTreeIteratorBase<proton::DocumentMetaStore::DocId, BTreeNoLeafData, NoAggregated, BTreeDefaultTraits::INTERNAL_SLOTS, BTreeDefaultTraits::LEAF_SLOTS, BTreeDefaultTraits::PATH_SIZE>;
+extern template class BTreeIteratorBase<proton::documentmetastore::GidToLidMapKey, BTreeNoLeafData, NoAggregated, BTreeDefaultTraits::INTERNAL_SLOTS, BTreeDefaultTraits::LEAF_SLOTS, BTreeDefaultTraits::PATH_SIZE>;
 
-extern template class BTreeConstIterator<proton::DocumentMetaStore::DocId, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
+extern template class BTreeConstIterator<proton::documentmetastore::GidToLidMapKey, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
 
-extern template class BTreeIterator<proton::DocumentMetaStore::DocId, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
+extern template class BTreeIterator<proton::documentmetastore::GidToLidMapKey, BTreeNoLeafData, NoAggregated, const proton::DocumentMetaStore::KeyComp &>;
 
 }
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.cpp b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.cpp
index 7d85c01ee74..94bba1bf64f 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.cpp
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.cpp
@@ -37,7 +37,8 @@ public:
           _writeDocSize(writeDocSize)
     { }
 
-    void operator()(uint32_t lid) {
+    void operator()(documentmetastore::GidToLidMapKey key) {
+        auto lid = key.get_lid();
         assert(lid < _metaDataStoreSize);
         const RawDocumentMetaData &metaData = _metaDataStore[lid];
         const GlobalId &gid = metaData.getGid();
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.h b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.h
index e8a6a7929f0..2f745929667 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.h
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/documentmetastoresaver.h
@@ -21,7 +21,7 @@ public:
     using KeyComp = documentmetastore::LidGidKeyComparator;
     using DocId = documentmetastore::IStore::DocId;
     using GidIterator = vespalib::btree::BTreeConstIterator<
-        DocId,
+        documentmetastore::GidToLidMapKey,
         vespalib::btree::BTreeNoLeafData,
         vespalib::btree::NoAggregated,
         const KeyComp &>;
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/gid_to_lid_map_key.cpp b/searchcore/src/vespa/searchcore/proton/documentmetastore/gid_to_lid_map_key.cpp
new file mode 100644
index 00000000000..eb31d6ed93a
--- /dev/null
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/gid_to_lid_map_key.cpp
@@ -0,0 +1,33 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "gid_to_lid_map_key.h"
+#include <vespa/document/base/globalid.h>
+
+using document::GlobalId;
+
+namespace proton::documentmetastore {
+
+GidToLidMapKey::GidToLidMapKey()
+    : _lid(FIND_DOC_ID),
+      _gid_key(0u)
+{
+}
+
+GidToLidMapKey::GidToLidMapKey(uint32_t lid, uint32_t gid_key)
+    : _lid(lid),
+      _gid_key(gid_key)
+{
+}
+
+GidToLidMapKey::GidToLidMapKey(uint32_t lid, const GlobalId& gid)
+    : GidToLidMapKey(lid, GlobalId::BucketOrderCmp::gid_key32(gid))
+{
+}
+
+GidToLidMapKey
+GidToLidMapKey::make_find_key(const GlobalId& gid)
+{
+    return GidToLidMapKey(FIND_DOC_ID, gid);
+}
+
+}
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/gid_to_lid_map_key.h b/searchcore/src/vespa/searchcore/proton/documentmetastore/gid_to_lid_map_key.h
new file mode 100644
index 00000000000..65e5665dbdf
--- /dev/null
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/gid_to_lid_map_key.h
@@ -0,0 +1,32 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <cstdint>
+#include <limits>
+
+namespace document { class GlobalId; }
+
+namespace proton::documentmetastore {
+
+/*
+ * Class containing lid and the most significant portion of gid according
+ * to compare functor (document::GlobalId::BucketOrderCmp).
+ */
+class GidToLidMapKey {
+    uint32_t _lid;
+    uint32_t _gid_key;
+    static constexpr uint32_t FIND_DOC_ID = std::numeric_limits<uint32_t>::max();
+
+public:
+    GidToLidMapKey();
+    GidToLidMapKey(uint32_t lid, uint32_t gid_key);
+    GidToLidMapKey(uint32_t lid, const document::GlobalId &gid);
+    static GidToLidMapKey make_find_key(const document::GlobalId &gid);
+
+    uint32_t get_lid() const { return _lid; }
+    uint32_t get_gid_key() const { return _gid_key; }
+    bool is_find_key() const { return _lid == FIND_DOC_ID; }
+};
+
+}
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/i_document_meta_store.h b/searchcore/src/vespa/searchcore/proton/documentmetastore/i_document_meta_store.h
index b757d2efb6a..0bb4c2913f3 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/i_document_meta_store.h
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/i_document_meta_store.h
@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include "gid_to_lid_map_key.h"
 #include "lid_gid_key_comparator.h"
 #include "i_simple_document_meta_store.h"
 #include <vespa/searchlib/attribute/attributeguard.h>
@@ -34,7 +35,7 @@ struct IDocumentMetaStore : public search::IDocumentMetaStore,
     // Lids are stored as keys in the tree, sorted by their gid counterpart.
     // The LidGidKeyComparator class maps from lids -> metadata by using the metadata store.
     // TODO(geirst): move this typedef and iterator functions away from this interface.
-    typedef vespalib::btree::BTree<DocId,
+    typedef vespalib::btree::BTree<documentmetastore::GidToLidMapKey,
             vespalib::btree::BTreeNoLeafData,
             vespalib::btree::NoAggregated,
             const documentmetastore::LidGidKeyComparator &> TreeType;
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.cpp b/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.cpp
index 48b90aa1baf..29caf1ae551 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.cpp
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.cpp
@@ -1,13 +1,9 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "lid_gid_key_comparator.h"
-#include <limits>
 
 namespace proton::documentmetastore {
 
-const search::IDocumentMetaStore::DocId
-LidGidKeyComparator::FIND_DOC_ID = std::numeric_limits<DocId>::max();
-
 LidGidKeyComparator::LidGidKeyComparator(const document::GlobalId &gid,
                                          const MetaDataStore &metaDataStore)
     : _gid(gid),
diff --git a/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.h b/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.h
index 4e1f2d72dad..a8a550c7a73 100644
--- a/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.h
+++ b/searchcore/src/vespa/searchcore/proton/documentmetastore/lid_gid_key_comparator.h
@@ -3,6 +3,7 @@
 #pragma once
 
 #include "raw_document_meta_data.h"
+#include "gid_to_lid_map_key.h"
 #include <vespa/document/base/globalid.h>
 #include <vespa/searchlib/common/idocumentmetastore.h>
 #include <vespa/vespalib/util/rcuvector.h>
@@ -15,9 +16,6 @@ namespace proton::documentmetastore {
  **/
 class LidGidKeyComparator
 {
-public:
-    static const search::IDocumentMetaStore::DocId FIND_DOC_ID;
-
 private:
     typedef search::IDocumentMetaStore::DocId DocId;
     typedef vespalib::RcuVectorBase<RawDocumentMetaData> MetaDataStore;
@@ -26,9 +24,9 @@ private:
     const MetaDataStore      &_metaDataStore;
     const document::GlobalId::BucketOrderCmp _gidCompare;
 
-    const document::GlobalId &getGid(DocId lid) const {
-        if (lid != FIND_DOC_ID) {
-            return _metaDataStore[lid].getGid();
+    const document::GlobalId &getGid(const GidToLidMapKey &key) const {
+        if (!key.is_find_key()) {
+            return _metaDataStore[key.get_lid()].getGid();
         }
         return _gid;
     }
@@ -36,7 +34,7 @@ private:
 public:
     /**
      * Creates a comparator that returns the given gid if
-     * FIND_DOC_ID is encountered. Otherwise the metadata store is
+     * key is a find key. Otherwise the metadata store is
      * used to map from lid -> metadata (including gid).
      **/
     LidGidKeyComparator(const document::GlobalId &gid,
@@ -45,8 +43,11 @@ public:
     LidGidKeyComparator(const RawDocumentMetaData &metaData,
                         const MetaDataStore &metaDataStore);
 
-    bool operator()(const DocId &lhs, const DocId &rhs) const {
-            return _gidCompare(getGid(lhs), getGid(rhs));
+    bool operator()(const GidToLidMapKey &lhs, const GidToLidMapKey &rhs) const {
+        if (lhs.get_gid_key() != rhs.get_gid_key()) {
+            return lhs.get_gid_key() < rhs.get_gid_key();
+        }
+        return _gidCompare(getGid(lhs), getGid(rhs));
     }
 
 };
diff --git a/searchcore/src/vespa/searchcore/proton/server/document_scan_iterator.cpp b/searchcore/src/vespa/searchcore/proton/server/document_scan_iterator.cpp
index 540e4d8785e..cbab4b72971 100644
--- a/searchcore/src/vespa/searchcore/proton/server/document_scan_iterator.cpp
+++ b/searchcore/src/vespa/searchcore/proton/server/document_scan_iterator.cpp
@@ -32,9 +32,9 @@ DocumentScanIterator::next(uint32_t compactLidLimit,
             (retry ? _metaStore.lowerBound(_lastGid) : _metaStore.upperBound(_lastGid))
                     : _metaStore.begin());
     uint32_t i = 1; // We have already 'scanned' a document when creating the iterator
-    for (; i < maxDocsToScan && itr.valid() && itr.getKey() <= compactLidLimit; ++i, ++itr) {}
+    for (; i < maxDocsToScan && itr.valid() && itr.getKey().get_lid() <= compactLidLimit; ++i, ++itr) {}
     if (itr.valid()) {
-        uint32_t lid = itr.getKey();
+        uint32_t lid = itr.getKey().get_lid();
         const RawDocumentMetaData &metaData = _metaStore.getRawMetaData(lid);
         _lastGid = metaData.getGid();
         _lastGidValid = true;
diff --git a/searchcore/src/vespa/searchcore/proton/server/documentbucketmover.cpp b/searchcore/src/vespa/searchcore/proton/server/documentbucketmover.cpp
index 9025fc56874..e10f7937dd0 100644
--- a/searchcore/src/vespa/searchcore/proton/server/documentbucketmover.cpp
+++ b/searchcore/src/vespa/searchcore/proton/server/documentbucketmover.cpp
@@ -116,7 +116,7 @@ DocumentBucketMover::moveDocuments(size_t maxDocsToMove)
     typedef std::vector<MoveKey> MoveVec;
     MoveVec toMove;
     for (; itr != end && docsMoved < maxDocsToMove; ++itr) {
-        DocumentIdT lid = itr.getKey();
+        DocumentIdT lid = itr.getKey().get_lid();
         const RawDocumentMetaData &metaData = _source->meta_store()->getRawMetaData(lid);
         if (metaData.getBucketUsedBits() != _bucket.getUsedBits()) {
             ++docsSkipped;