Merge pull request #28556 from vespa-engine/toregge/add-comparator-to-unique-store

Add comparator to unique store.

5 files changed, 53 insertions, 59 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/enumstore.h b/searchlib/src/vespa/searchlib/attribute/enumstore.h
index 0b16ec3acc1..8cafd04586e 100644
--- a/searchlib/src/vespa/searchlib/attribute/enumstore.h
+++ b/searchlib/src/vespa/searchlib/attribute/enumstore.h
@@ -55,7 +55,6 @@ private:
     UniqueStoreType        _store;
     IEnumStoreDictionary*  _dict;
     bool                   _is_folded;
-    ComparatorType         _comparator;
     ComparatorType         _foldedComparator;
     enumstore::EnumStoreCompactionSpec _compaction_spec;
     EntryType              _default_value;
@@ -187,8 +186,8 @@ public:
         return BatchUpdater(*this);
     }
 
-    const EntryComparator & get_comparator() const {
-        return _comparator;
+    const EntryComparator & get_comparator() const noexcept {
+        return _store.get_comparator();
     }
 
     ComparatorType make_comparator(const EntryType& lookup_value) const {
diff --git a/searchlib/src/vespa/searchlib/attribute/enumstore.hpp b/searchlib/src/vespa/searchlib/attribute/enumstore.hpp
index c0eebee8e94..85cf676182e 100644
--- a/searchlib/src/vespa/searchlib/attribute/enumstore.hpp
+++ b/searchlib/src/vespa/searchlib/attribute/enumstore.hpp
@@ -77,7 +77,6 @@ EnumStoreT<EntryT>::EnumStoreT(bool has_postings, const DictionaryConfig& dict_c
     : _store(std::move(memory_allocator)),
       _dict(),
       _is_folded(dict_cfg.getMatch() == DictionaryConfig::Match::UNCASED),
-      _comparator(_store.get_data_store()),
       _foldedComparator(make_optionally_folded_comparator(is_folded())),
       _compaction_spec(),
       _default_value(default_value),
diff --git a/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp b/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp
index a9826a71091..27e2c274dc6 100644
--- a/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp
+++ b/vespalib/src/tests/datastore/unique_store/unique_store_test.cpp
@@ -39,7 +39,7 @@ struct TestBase : public ::testing::Test {
     using EntryRefType = typename UniqueStoreType::RefType;
     using ValueType = typename UniqueStoreType::EntryType;
     using ValueConstRefType = typename UniqueStoreType::EntryConstRefType;
-    using CompareType = typename UniqueStoreType::CompareType;
+    using ComparatorType = typename UniqueStoreType::ComparatorType;
     using ReferenceStoreValueType = std::conditional_t<std::is_same_v<ValueType, const char *>, std::string, ValueType>;
     using ReferenceStore = std::map<EntryRef, std::pair<ReferenceStoreValueType,uint32_t>>;
 
@@ -164,13 +164,13 @@ TestBase<UniqueStoreTypeAndDictionaryType>::TestBase()
         EXPECT_FALSE(store.get_dictionary().get_has_hash_dictionary());
         break;
     case DictionaryType::BTREE_AND_HASH:
-        store.set_dictionary(std::make_unique<UniqueStoreDictionary<uniquestore::DefaultDictionary, IUniqueStoreDictionary, ShardedHashMap>>(std::make_unique<CompareType>(store.get_data_store())));
+        store.set_dictionary(std::make_unique<UniqueStoreDictionary<uniquestore::DefaultDictionary, IUniqueStoreDictionary, ShardedHashMap>>(std::make_unique<ComparatorType>(store.get_data_store())));
         EXPECT_TRUE(store.get_dictionary().get_has_btree_dictionary());
         EXPECT_TRUE(store.get_dictionary().get_has_hash_dictionary());
         break;
     case DictionaryType::HASH:
     default:
-        store.set_dictionary(std::make_unique<UniqueStoreDictionary<NoBTreeDictionary, IUniqueStoreDictionary, ShardedHashMap>>(std::make_unique<CompareType>(store.get_data_store())));
+        store.set_dictionary(std::make_unique<UniqueStoreDictionary<NoBTreeDictionary, IUniqueStoreDictionary, ShardedHashMap>>(std::make_unique<ComparatorType>(store.get_data_store())));
         EXPECT_FALSE(store.get_dictionary().get_has_btree_dictionary());
         EXPECT_TRUE(store.get_dictionary().get_has_hash_dictionary());
     }
@@ -464,7 +464,7 @@ TEST_F(DoubleTest, nan_is_handled)
 
 TEST_F(DoubleTest, control_memory_usage) {
     static constexpr size_t sizeof_deque = vespalib::datastore::DataStoreBase::sizeof_entry_ref_hold_list_deque;
-    EXPECT_EQ(376u + sizeof_deque, sizeof(store));
+    EXPECT_EQ(400u + sizeof_deque, sizeof(store));
     EXPECT_EQ(120u, sizeof(BufferState));
     EXPECT_EQ(28740u, store.get_values_memory_usage().allocatedBytes());
     EXPECT_EQ(24780u, store.get_values_memory_usage().usedBytes());
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store.h b/vespalib/src/vespa/vespalib/datastore/unique_store.h
index f2d62020ab4..9589141e350 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store.h
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store.h
@@ -12,6 +12,7 @@
 #include "unique_store_allocator.h"
 #include "unique_store_comparator.h"
 #include "unique_store_entry.h"
+#include <functional>
 
 namespace vespalib::alloc { class MemoryAllocator; }
 
@@ -30,14 +31,14 @@ class UniqueStoreRemapper;
  * Datastore for unique values of type EntryT that is accessed via a
  * 32-bit EntryRef.
  */
-template <typename EntryT, typename RefT = EntryRefT<22>, typename Compare = UniqueStoreComparator<EntryT, RefT>, typename Allocator = UniqueStoreAllocator<EntryT, RefT> >
+template <typename EntryT, typename RefT = EntryRefT<22>, typename Comparator = UniqueStoreComparator<EntryT, RefT>, typename Allocator = UniqueStoreAllocator<EntryT, RefT> >
 class UniqueStore
 {
 public:
     using DataStoreType = DataStoreT<RefT>;
     using EntryType = EntryT;
     using RefType = RefT;
-    using CompareType = Compare;
+    using ComparatorType = Comparator;
     using Enumerator = UniqueStoreEnumerator<RefT>;
     using Builder = UniqueStoreBuilder<Allocator>;
     using Remapper = UniqueStoreRemapper<RefT>;
@@ -45,12 +46,12 @@ public:
 private:
     Allocator _allocator;
     DataStoreType &_store;
+    ComparatorType _comparator;
     std::unique_ptr<IUniqueStoreDictionary> _dict;
     using generation_t = vespalib::GenerationHandler::generation_t;
 
 public:
-    UniqueStore(std::shared_ptr<alloc::MemoryAllocator> memory_allocator);
-    UniqueStore(std::unique_ptr<IUniqueStoreDictionary> dict, std::shared_ptr<alloc::MemoryAllocator> memory_allocator);
+    UniqueStore(std::shared_ptr<alloc::MemoryAllocator> memory_allocator, const std::function<ComparatorType(const DataStoreType&)>& comparator_factory = [](const auto& data_store) { return ComparatorType(data_store);});
     ~UniqueStore();
     void set_dictionary(std::unique_ptr<IUniqueStoreDictionary> dict);
     UniqueStoreAddResult add(EntryConstRefType value);
@@ -81,6 +82,8 @@ public:
     Builder getBuilder(uint32_t uniqueValuesHint);
     Enumerator getEnumerator(bool sort_unique_values);
 
+    const ComparatorType& get_comparator() const noexcept { return _comparator; }
+
     // Should only be used for unit testing
     const BufferState &bufferState(EntryRef ref) const;
 };
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
index 0efaf04b26e..fc912382cf8 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store.hpp
@@ -27,60 +27,53 @@ using DefaultUniqueStoreDictionary = UniqueStoreDictionary<DefaultDictionary>;
 
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
-UniqueStore<EntryT, RefT, Compare, Allocator>::UniqueStore(std::shared_ptr<alloc::MemoryAllocator> memory_allocator)
-    : UniqueStore<EntryT, RefT, Compare, Allocator>(std::make_unique<uniquestore::DefaultUniqueStoreDictionary>(std::unique_ptr<EntryComparator>()), std::move(memory_allocator))
-{
-}
-
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
-UniqueStore<EntryT, RefT, Compare, Allocator>::UniqueStore(std::unique_ptr<IUniqueStoreDictionary> dict, std::shared_ptr<alloc::MemoryAllocator> memory_allocator)
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
+UniqueStore<EntryT, RefT, Comparator, Allocator>::UniqueStore(std::shared_ptr<alloc::MemoryAllocator> memory_allocator, const std::function<ComparatorType(const DataStoreType&)>& comparator_factory)
     : _allocator(std::move(memory_allocator)),
       _store(_allocator.get_data_store()),
-      _dict(std::move(dict))
+      _comparator(comparator_factory(_store)),
+      _dict(std::make_unique<uniquestore::DefaultUniqueStoreDictionary>(std::unique_ptr<EntryComparator>()))
 {
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
-UniqueStore<EntryT, RefT, Compare, Allocator>::~UniqueStore() = default;
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
+UniqueStore<EntryT, RefT, Comparator, Allocator>::~UniqueStore() = default;
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 void
-UniqueStore<EntryT, RefT, Compare, Allocator>::set_dictionary(std::unique_ptr<IUniqueStoreDictionary> dict)
+UniqueStore<EntryT, RefT, Comparator, Allocator>::set_dictionary(std::unique_ptr<IUniqueStoreDictionary> dict)
 {
     _dict = std::move(dict);
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 UniqueStoreAddResult
-UniqueStore<EntryT, RefT, Compare, Allocator>::add(EntryConstRefType value)
+UniqueStore<EntryT, RefT, Comparator, Allocator>::add(EntryConstRefType value)
 {
-    Compare comp(_store, value);
+    Comparator comp(_store, value);
     UniqueStoreAddResult result = _dict->add(comp, [this, &value]() -> EntryRef { return _allocator.allocate(value); });
     _allocator.get_wrapped(result.ref()).inc_ref_count();
     return result;
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 EntryRef
-UniqueStore<EntryT, RefT, Compare, Allocator>::find(EntryConstRefType value)
+UniqueStore<EntryT, RefT, Comparator, Allocator>::find(EntryConstRefType value)
 {
-    Compare comp(_store, value);
+    Comparator comp(_store, value);
     return _dict->find(comp);
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 void
-UniqueStore<EntryT, RefT, Compare, Allocator>::remove(EntryRef ref)
+UniqueStore<EntryT, RefT, Comparator, Allocator>::remove(EntryRef ref)
 {
     auto &wrapped_entry = _allocator.get_wrapped(ref);
     auto ref_count = wrapped_entry.get_ref_count();
     assert(ref_count > 0u);
     wrapped_entry.dec_ref_count();
     if (ref_count == 1u) {
-        EntryType unused{};
-        Compare comp(_store, unused);
-        _dict->remove(comp, ref);
+        _dict->remove(_comparator, ref);
         _allocator.hold(ref);
     }
 }
@@ -152,9 +145,9 @@ public:
 
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
-std::unique_ptr<typename UniqueStore<EntryT, RefT, Compare, Allocator>::Remapper>
-UniqueStore<EntryT, RefT, Compare, Allocator>::compact_worst(CompactionSpec compaction_spec, const CompactionStrategy& compaction_strategy)
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
+std::unique_ptr<typename UniqueStore<EntryT, RefT, Comparator, Allocator>::Remapper>
+UniqueStore<EntryT, RefT, Comparator, Allocator>::compact_worst(CompactionSpec compaction_spec, const CompactionStrategy& compaction_strategy)
 {
     auto compacting_buffers = _store.start_compact_worst_buffers(compaction_spec, compaction_strategy);
     if (compacting_buffers->empty()) {
@@ -164,71 +157,71 @@ UniqueStore<EntryT, RefT, Compare, Allocator>::compact_worst(CompactionSpec comp
     }
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 vespalib::MemoryUsage
-UniqueStore<EntryT, RefT, Compare, Allocator>::getMemoryUsage() const
+UniqueStore<EntryT, RefT, Comparator, Allocator>::getMemoryUsage() const
 {
     vespalib::MemoryUsage usage = get_values_memory_usage();
     usage.merge(get_dictionary_memory_usage());
     return usage;
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 vespalib::AddressSpace
-UniqueStore<EntryT, RefT, Compare, Allocator>::get_values_address_space_usage() const
+UniqueStore<EntryT, RefT, Comparator, Allocator>::get_values_address_space_usage() const
 {
     return _allocator.get_data_store().getAddressSpaceUsage();
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 const BufferState &
-UniqueStore<EntryT, RefT, Compare, Allocator>::bufferState(EntryRef ref) const
+UniqueStore<EntryT, RefT, Comparator, Allocator>::bufferState(EntryRef ref) const
 {
     RefT internalRef(ref);
     return _store.getBufferState(internalRef.bufferId());
 }
 
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 void
-UniqueStore<EntryT, RefT, Compare, Allocator>::assign_generation(generation_t current_gen)
+UniqueStore<EntryT, RefT, Comparator, Allocator>::assign_generation(generation_t current_gen)
 {
     _dict->assign_generation(current_gen);
     _store.assign_generation(current_gen);
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 void
-UniqueStore<EntryT, RefT, Compare, Allocator>::reclaim_memory(generation_t oldest_used_gen)
+UniqueStore<EntryT, RefT, Comparator, Allocator>::reclaim_memory(generation_t oldest_used_gen)
 {
     _dict->reclaim_memory(oldest_used_gen);
     _store.reclaim_memory(oldest_used_gen);
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 void
-UniqueStore<EntryT, RefT, Compare, Allocator>::freeze()
+UniqueStore<EntryT, RefT, Comparator, Allocator>::freeze()
 {
     _dict->freeze();
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
-typename UniqueStore<EntryT, RefT, Compare, Allocator>::Builder
-UniqueStore<EntryT, RefT, Compare, Allocator>::getBuilder(uint32_t uniqueValuesHint)
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
+typename UniqueStore<EntryT, RefT, Comparator, Allocator>::Builder
+UniqueStore<EntryT, RefT, Comparator, Allocator>::getBuilder(uint32_t uniqueValuesHint)
 {
     return Builder(_allocator, *_dict, uniqueValuesHint);
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
-typename UniqueStore<EntryT, RefT, Compare, Allocator>::Enumerator
-UniqueStore<EntryT, RefT, Compare, Allocator>::getEnumerator(bool sort_unique_values)
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
+typename UniqueStore<EntryT, RefT, Comparator, Allocator>::Enumerator
+UniqueStore<EntryT, RefT, Comparator, Allocator>::getEnumerator(bool sort_unique_values)
 {
     return Enumerator(*_dict, _store, sort_unique_values);
 }
 
-template <typename EntryT, typename RefT, typename Compare, typename Allocator>
+template <typename EntryT, typename RefT, typename Comparator, typename Allocator>
 uint32_t
-UniqueStore<EntryT, RefT, Compare, Allocator>::getNumUniques() const
+UniqueStore<EntryT, RefT, Comparator, Allocator>::getNumUniques() const
 {
     return _dict->get_num_uniques();
 }