Stop requiring dictionary iterator when handling diversify over a single

dictionary entry.

3 files changed, 28 insertions, 15 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/diversity.h b/searchlib/src/vespa/searchlib/attribute/diversity.h
index 6e3c6cfff00..ff7d9b5f83d 100644
--- a/searchlib/src/vespa/searchlib/attribute/diversity.h
+++ b/searchlib/src/vespa/searchlib/attribute/diversity.h
@@ -123,4 +123,22 @@ void diversify(bool forward, const DictItr &lower, const DictItr &upper, const P
     }
 }
 
+template <typename PostingStore, typename Result>
+void diversify_single(vespalib::datastore::EntryRef posting_idx, const PostingStore &posting, size_t wanted_hits,
+               const IAttributeVector &diversity_attr, size_t max_per_group,
+               size_t cutoff_max_groups, bool cutoff_strict,
+               Result &result, std::vector<size_t> &fragments)
+{
+    auto filter = DiversityFilter::create(diversity_attr, wanted_hits, max_per_group, cutoff_max_groups, cutoff_strict);
+    DiversityRecorder<Result> recorder(*filter, result);
+    using DataType = typename PostingStore::DataType;
+    using KeyDataType = typename PostingStore::KeyDataType;
+    posting.foreach_frozen(posting_idx,
+                           [&](uint32_t key, const DataType &data)
+                           { recorder.push_back(KeyDataType(key, data)); });
+    if (fragments.back() < result.size()) {
+        fragments.push_back(result.size());
+    }
+}
+
 }
diff --git a/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.cpp b/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.cpp
index c7c0805b1a4..08db7b2b6b7 100644
--- a/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.cpp
@@ -45,19 +45,9 @@ PostingListSearchContext::~PostingListSearchContext() = default;
 void
 PostingListSearchContext::lookupTerm(const vespalib::datastore::EntryComparator &comp)
 {
-    if (_dictionary.get_has_unordered_dictionary()) {
-        // Note: Diversity requires _lowerDictItr and upperDictItr to be setup
-        auto lookup_result = _dictionary.find_posting_list(comp, _frozenDictionary.getRoot());
-        if (lookup_result.first.valid()) {
-            _pidx = lookup_result.second;
-            _uniqueValues = 1u;
-        }
-        return;
-    }
-    _lowerDictItr.lower_bound(_frozenDictionary.getRoot(), EnumIndex(), comp);
-    _upperDictItr = _lowerDictItr;
-    if (_upperDictItr.valid() && !comp.less(EnumIndex(), _upperDictItr.getKey())) {
-        ++_upperDictItr;
+    auto lookup_result = _dictionary.find_posting_list(comp, _frozenDictionary.getRoot());
+    if (lookup_result.first.valid()) {
+        _pidx = lookup_result.second;
         _uniqueValues = 1u;
     }
 }
diff --git a/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.hpp b/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.hpp
index f99ca61cba0..fb2617064e4 100644
--- a/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.hpp
+++ b/searchlib/src/vespa/searchlib/attribute/postinglistsearchcontext.hpp
@@ -138,8 +138,13 @@ PostingListSearchContextT<DataT>::diversify(bool forward, size_t wanted_hits, co
 {
     if (!_merger.merge_done()) {
         _merger.reserveArray(128, wanted_hits);
-        diversity::diversify(forward, _lowerDictItr, _upperDictItr, _postingList, wanted_hits, diversity_attr,
-                             max_per_group, cutoff_groups, cutoff_strict, _merger.getWritableArray(), _merger.getWritableStartPos());
+        if (_uniqueValues == 1u && !_lowerDictItr.valid() && _pidx.valid()) {
+            diversity::diversify_single(_pidx, _postingList, wanted_hits, diversity_attr,
+                                        max_per_group, cutoff_groups, cutoff_strict, _merger.getWritableArray(), _merger.getWritableStartPos());
+        } else {
+            diversity::diversify(forward, _lowerDictItr, _upperDictItr, _postingList, wanted_hits, diversity_attr,
+                                 max_per_group, cutoff_groups, cutoff_strict, _merger.getWritableArray(), _merger.getWritableStartPos());
+        }
         _merger.merge();
     }
 }