Merge pull request #3123 from vespa-engine/toregge/maintain-direct-forward-mapping-in-reference-mappings

Maintain direct forward mapping from lid to referenced lid in reference

6 files changed, 81 insertions, 19 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp b/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp
index 3cef9a03c8f..c431b956c1d 100644
--- a/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp
@@ -18,7 +18,8 @@ ImportedSearchContext::ImportedSearchContext(
     : _imported_attribute(imported_attribute),
       _reference_attribute(*_imported_attribute.getReferenceAttribute()),
       _target_attribute(*_imported_attribute.getTargetAttribute()),
-      _target_search_context(_target_attribute.getSearch(std::move(term), params))
+      _target_search_context(_target_attribute.getSearch(std::move(term), params)),
+      _referencedLids(_reference_attribute.getReferencedLids())
 {
 }
 
@@ -60,11 +61,11 @@ const vespalib::string& ImportedSearchContext::attributeName() const {
 }
 
 bool ImportedSearchContext::cmp(DocId docId, int32_t& weight) const {
-    return _target_search_context->cmp(_reference_attribute.getReferencedLid(docId), weight);
+    return _target_search_context->cmp(_referencedLids[docId], weight);
 }
 
 bool ImportedSearchContext::cmp(DocId docId) const {
-    return _target_search_context->cmp(_reference_attribute.getReferencedLid(docId));
+    return _target_search_context->cmp(_referencedLids[docId]);
 }
 
 } // attribute
diff --git a/searchlib/src/vespa/searchlib/attribute/imported_search_context.h b/searchlib/src/vespa/searchlib/attribute/imported_search_context.h
index 3863297afb4..ce6642cd93f 100644
--- a/searchlib/src/vespa/searchlib/attribute/imported_search_context.h
+++ b/searchlib/src/vespa/searchlib/attribute/imported_search_context.h
@@ -4,6 +4,7 @@
 
 #include "attributevector.h"
 #include <vespa/searchcommon/attribute/i_search_context.h>
+#include <vespa/vespalib/util/arrayref.h>
 #include <memory>
 
 namespace search {
@@ -27,10 +28,12 @@ class SearchContextParams;
  * considered a match.
  */
 class ImportedSearchContext : public ISearchContext {
+    using ReferencedLids = vespalib::ConstArrayRef<uint32_t>;
     const ImportedAttributeVector&                  _imported_attribute;
     const ReferenceAttribute&                       _reference_attribute;
     const AttributeVector&                          _target_attribute;
     std::unique_ptr<AttributeVector::SearchContext> _target_search_context;
+    ReferencedLids                                  _referencedLids;
 public:
     ImportedSearchContext(std::unique_ptr<QueryTermSimple> term,
                           const SearchContextParams& params,
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp b/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
index 480a9178655..147c656ccfd 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
@@ -58,6 +58,7 @@ void
 ReferenceAttribute::onAddDocs(DocId limit)
 {
     _indices.reserve(limit);
+    _referenceMappings.onAddDocs(limit);
 }
 
 bool
@@ -66,6 +67,7 @@ ReferenceAttribute::addDoc(DocId &doc)
     bool incGen = _indices.isFull();
     doc = _indices.size();
     _indices.push_back(EntryRef());
+    _referenceMappings.addDoc();
     incNumDocs();
     updateUncommittedDocIdLimit(doc);
     if (incGen) {
@@ -221,6 +223,7 @@ ReferenceAttribute::onLoad()
         builder.add(value);
     }
     builder.setupRefCounts();
+    _referenceMappings.onLoad(numDocs);
     _indices.clear();
     _indices.unsafe_reserve(numDocs);
     for (uint32_t doc = 0; doc < numDocs; ++doc) {
@@ -314,18 +317,6 @@ ReferenceAttribute::setGidToLidMapperFactory(std::shared_ptr<IGidToLidMapperFact
     _gidToLidMapperFactory = std::move(gidToLidMapperFactory);
 }
 
-ReferenceAttribute::DocId
-ReferenceAttribute::getReferencedLid(DocId doc) const
-{
-    assert(doc < _indices.size());
-    EntryRef ref = _indices[doc];
-    if (!ref.valid()) {
-        return 0;
-    } else {
-        return _store.get(ref).lid();
-    }
-}
-
 void
 ReferenceAttribute::notifyGidToLidChange(const GlobalId &gid, DocId referencedLid)
 {
@@ -348,7 +339,8 @@ ReferenceAttribute::populateReferencedLids()
         saver.foreach_key([&store,&mapper,this](EntryRef ref)
                           {   const Reference &entry = store.get(ref);
                               entry.setLid(mapper.mapGidToLid(entry.gid()));
-                              _referenceMappings.syncReverseMappingIndices(entry); });
+                              _referenceMappings.syncMappings(entry);
+                          });
     }
     commit();
 }
@@ -375,6 +367,7 @@ ReferenceAttribute::onShrinkLidSpace()
     uint32_t committedDocIdLimit = getCommittedDocIdLimit();
     assert(_indices.size() >= committedDocIdLimit);
     _indices.shrink(committedDocIdLimit);
+    _referenceMappings.shrink(committedDocIdLimit);
     setNumDocs(committedDocIdLimit);
 }
 
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.h b/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
index 08bcfc2cbb8..53b70c1da1d 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
+++ b/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
@@ -35,6 +35,7 @@ public:
                                              std::less<uint32_t>,
                                              btree::BTreeDefaultTraits,
                                              btree::NoAggrCalc>;
+    using ReferencedLids = ReferenceMappings::ReferencedLids;
 private:
     ReferenceStore _store;
     ReferenceStoreIndices _indices;
@@ -71,7 +72,8 @@ public:
     const Reference *getReference(DocId doc);
     void setGidToLidMapperFactory(std::shared_ptr<IGidToLidMapperFactory> gidToLidMapperFactory);
     std::shared_ptr<IGidToLidMapperFactory> getGidToLidMapperFactory() const { return _gidToLidMapperFactory; }
-    DocId getReferencedLid(DocId doc) const;
+    ReferencedLids getReferencedLids() const { return _referenceMappings.getReferencedLids(); }
+    DocId getReferencedLid(DocId doc) const { return _referenceMappings.getReferencedLid(doc); }
     void notifyGidToLidChange(const GlobalId &gid, DocId referencedLid);
     void populateReferencedLids();
     virtual void clearDocs(DocId lidLow, DocId lidLimit) override;
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp b/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp
index 67f02c50d88..175dc05396d 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp
@@ -9,7 +9,8 @@ namespace search::attribute {
 
 ReferenceMappings::ReferenceMappings(GenerationHolder &genHolder)
     : _reverseMappingIndices(genHolder),
-      _reverseMapping()
+      _reverseMapping(),
+      _referencedLids(genHolder)
 {
 }
 
@@ -27,6 +28,18 @@ ReferenceMappings::clearMapping(const Reference &entry)
 }
 
 void
+ReferenceMappings::syncForwardMapping(const Reference &entry)
+{
+    uint32_t referencedLid = entry.lid();
+    EntryRef revMapIdx = entry.revMapIdx();
+    auto &referencedLids = _referencedLids;
+    _reverseMapping.foreach_unfrozen_key(revMapIdx,
+                                         [&referencedLids, referencedLid](uint32_t lid)
+                                         { referencedLids[lid] = referencedLid; });
+}
+
+
+void
 ReferenceMappings::syncReverseMappingIndices(const Reference &entry)
 {
     uint32_t referencedLid = entry.lid();
@@ -37,6 +50,13 @@ ReferenceMappings::syncReverseMappingIndices(const Reference &entry)
 }
 
 void
+ReferenceMappings::syncMappings(const Reference &entry)
+{
+    syncReverseMappingIndices(entry);
+    syncForwardMapping(entry);
+}
+
+void
 ReferenceMappings::removeReverseMapping(const Reference &entry, uint32_t lid)
 {
     EntryRef revMapIdx = entry.revMapIdx();
@@ -44,6 +64,7 @@ ReferenceMappings::removeReverseMapping(const Reference &entry, uint32_t lid)
     std::atomic_thread_fence(std::memory_order_release);
     entry.setRevMapIdx(revMapIdx);
     syncReverseMappingIndices(entry);
+    _referencedLids[lid] = 0; // forward mapping
 }
 
 void
@@ -55,6 +76,7 @@ ReferenceMappings::addReverseMapping(const Reference &entry, uint32_t lid)
     std::atomic_thread_fence(std::memory_order_release);
     entry.setRevMapIdx(revMapIdx);
     syncReverseMappingIndices(entry);
+    _referencedLids[lid] = entry.lid(); // forward mapping
 }
 
 void
@@ -77,6 +99,32 @@ ReferenceMappings::notifyGidToLidChange(const Reference &entry, uint32_t referen
         entry.setLid(referencedLid);
     }
     syncReverseMappingIndices(entry);
+    syncForwardMapping(entry);
+}
+
+void
+ReferenceMappings::onAddDocs(uint32_t docIdLimit)
+{
+    _referencedLids.reserve(docIdLimit);
+}
+
+void
+ReferenceMappings::addDoc()
+{
+    _referencedLids.push_back(0);
+}
+
+void
+ReferenceMappings::onLoad(uint32_t docIdLimit)
+{
+    _referencedLids.clear();
+    _referencedLids.unsafe_reserve(docIdLimit);
+}
+
+void
+ReferenceMappings::shrink(uint32_t docIdLimit)
+{
+    _referencedLids.shrink(docIdLimit);
 }
 
 }
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_mappings.h b/searchlib/src/vespa/searchlib/attribute/reference_mappings.h
index 129ccbbea25..9880bf44efc 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_mappings.h
+++ b/searchlib/src/vespa/searchlib/attribute/reference_mappings.h
@@ -30,9 +30,15 @@ class ReferenceMappings
     // Store of B-Trees, used to map from gid or referenced lid to
     // referencing lids.
     ReverseMapping _reverseMapping;
+    // vector containing referenced lid given referencing lid
+    RcuVectorBase<uint32_t> _referencedLids;
 
+    void syncForwardMapping(const Reference &entry);
+    void syncReverseMappingIndices(const Reference &entry);
 
 public:
+    using ReferencedLids = vespalib::ConstArrayRef<uint32_t>;
+
     ReferenceMappings(GenerationHolder &genHolder);
 
     ~ReferenceMappings();
@@ -48,9 +54,15 @@ public:
 
     // Handle mapping changes
     void notifyGidToLidChange(const Reference &entry, uint32_t referencedLid);
-    void syncReverseMappingIndices(const Reference &entry);
     void removeReverseMapping(const Reference &entry, uint32_t lid);
     void addReverseMapping(const Reference &entry, uint32_t lid);
+    void syncMappings(const Reference &entry);
+
+    // Maintain size of mapping from lid to referenced lid
+    void onAddDocs(uint32_t docIdLimit);
+    void addDoc();
+    void onLoad(uint32_t docIdLimit);
+    void shrink(uint32_t docIdLimit);
 
     // Setup mapping after load
     void buildReverseMapping(const Reference &entry, const std::vector<ReverseMapping::KeyDataType> &adds);
@@ -59,6 +71,9 @@ public:
     template <typename FunctionType>
     void
     foreach_lid(uint32_t referencedLid, FunctionType &&func) const;
+
+    ReferencedLids getReferencedLids() const { return ReferencedLids(&_referencedLids[0], _referencedLids.size()); }
+    uint32_t getReferencedLid(uint32_t doc) const { return _referencedLids[doc]; }
 };
 
 template <typename FunctionType>