allow profiling weak and children

- new decoration strategy: transform_children
- track child reordering
  (to make profile results match the structure of the blueprint tree)

7 files changed, 70 insertions, 1 deletions

diff --git a/searchlib/src/tests/queryeval/profiled_iterator/profiled_iterator_test.cpp b/searchlib/src/tests/queryeval/profiled_iterator/profiled_iterator_test.cpp
index d0942e14f7c..dbf01fac4f7 100644
--- a/searchlib/src/tests/queryeval/profiled_iterator/profiled_iterator_test.cpp
+++ b/searchlib/src/tests/queryeval/profiled_iterator/profiled_iterator_test.cpp
@@ -6,6 +6,8 @@
 #include <vespa/vespalib/util/require.h>
 #include <vespa/vespalib/data/slime/slime.h>
 #include <vespa/searchlib/queryeval/profiled_iterator.h>
+#include <vespa/searchlib/queryeval/wand/weak_and_heap.h>
+#include <vespa/searchlib/queryeval/wand/weak_and_search.h>
 #include <vespa/searchlib/queryeval/simplesearch.h>
 #include <vespa/searchlib/queryeval/sourceblendersearch.h>
 #include <vespa/searchlib/queryeval/andsearch.h>
@@ -84,6 +86,19 @@ SearchIterator::UP create_iterator_tree() {
                      t({2,4,6,8}), 5));
 }
 
+SearchIterator::UP create_weak_and() {
+    struct DummyHeap : WeakAndHeap {
+        void adjust(score_t *, score_t *) override {}
+        DummyHeap() : WeakAndHeap(100) {}
+    };
+    static DummyHeap dummy_heap;
+    WeakAndSearch::Terms terms;
+    terms.emplace_back(T({1,2,3}).release(), 100, 3);
+    terms.emplace_back(T({5,6}).release(), 200, 2);
+    terms.emplace_back(T({8}).release(), 300, 1);
+    return WeakAndSearch::create(terms, wand::MatchParams(dummy_heap), wand::TermFrequencyScorer(), 100, true, true);
+}
+
 void collect(std::map<vespalib::string,size_t> &counts, const auto &node) {
     if (!node.valid()) {
         return;
@@ -190,4 +205,25 @@ TEST(ProfiledIteratorTest, iterator_tree_can_be_profiled) {
     EXPECT_EQ(counts["/1/1/SimpleSearch/init"], 2);
 }
 
+TEST(ProfiledIteratorTest, weak_and_can_be_profiled) {
+    ExecutionProfiler profiler(64);
+    auto root = create_weak_and();
+    root = ProfiledIterator::profile(profiler, std::move(root));
+    fprintf(stderr, "%s", root->asString().c_str());
+    verify_result(*root, {1,2,3,5,6,8});
+    Slime slime;
+    profiler.report(slime.setObject());
+    fprintf(stderr, "%s", slime.toString().c_str());
+    auto counts = collect_counts(slime.get());
+    print_counts(counts);
+    EXPECT_EQ(counts["/WeakAndSearchLR/init"], 1);
+    EXPECT_EQ(counts["/0/SimpleSearch/init"], 1);
+    EXPECT_EQ(counts["/1/SimpleSearch/init"], 1);
+    EXPECT_EQ(counts["/2/SimpleSearch/init"], 1);
+    EXPECT_EQ(counts["/WeakAndSearchLR/seek"], 7);
+    EXPECT_EQ(counts["/0/SimpleSearch/seek"], 4);
+    EXPECT_EQ(counts["/1/SimpleSearch/seek"], 3);
+    EXPECT_EQ(counts["/2/SimpleSearch/seek"], 2);
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchlib/src/vespa/searchlib/queryeval/iterator_pack.h b/searchlib/src/vespa/searchlib/queryeval/iterator_pack.h
index f05bf8e1adc..7c4e874601e 100644
--- a/searchlib/src/vespa/searchlib/queryeval/iterator_pack.h
+++ b/searchlib/src/vespa/searchlib/queryeval/iterator_pack.h
@@ -61,6 +61,11 @@ public:
     }
     std::unique_ptr<BitVector> get_hits(uint32_t begin_id, uint32_t end_id) const;
     void or_hits_into(BitVector &result, uint32_t begin_id) const;
+    void transform_children(auto f) {
+        for (size_t i = 0; i < _children.size(); ++i) {
+            _children[i] = f(std::move(_children[i]), i);
+        }
+    }
 };
 
 using SearchIteratorPack = SearchIteratorPackT<uint16_t>;
diff --git a/searchlib/src/vespa/searchlib/queryeval/profiled_iterator.cpp b/searchlib/src/vespa/searchlib/queryeval/profiled_iterator.cpp
index ae0cc09b3ab..44b8d40ad1d 100644
--- a/searchlib/src/vespa/searchlib/queryeval/profiled_iterator.cpp
+++ b/searchlib/src/vespa/searchlib/queryeval/profiled_iterator.cpp
@@ -6,6 +6,7 @@
 #include <vespa/vespalib/objects/visit.hpp>
 #include <vespa/vespalib/util/classname.h>
 #include <vespa/vespalib/util/stringfmt.h>
+#include "wand/weak_and_search.h"
 
 #include <typeindex>
 
@@ -54,6 +55,10 @@ void handle_leaf_node(Profiler &profiler, SearchIterator &leaf, const vespalib::
             child = ProfiledIterator::profile(profiler, std::move(child), fmt("%s%zu/", path.c_str(), i));
             source_blender.setChild(i, std::move(child));
         }
+    } else if (auto *weak_and = leaf.as_weak_and(); weak_and != nullptr) {
+        weak_and->transform_children([&](auto child, size_t i){
+                                         return ProfiledIterator::profile(profiler, std::move(child), fmt("%s%zu/", path.c_str(), i));
+                                     });
     }
 }
 
diff --git a/searchlib/src/vespa/searchlib/queryeval/searchiterator.h b/searchlib/src/vespa/searchlib/queryeval/searchiterator.h
index f8124a161e2..14ca599232a 100644
--- a/searchlib/src/vespa/searchlib/queryeval/searchiterator.h
+++ b/searchlib/src/vespa/searchlib/queryeval/searchiterator.h
@@ -20,6 +20,8 @@ namespace search::attribute { class ISearchContext; }
 
 namespace search::queryeval {
 
+struct WeakAndSearch;
+
 /**
  * This is the abstract superclass of all search objects. Each search
  * object act as an iterator over documents that are results for the
@@ -357,6 +359,8 @@ public:
      */
     virtual bool isMultiSearch() const { return false; }
 
+    virtual WeakAndSearch *as_weak_and() noexcept { return nullptr; }
+
     /**
      * This is used for adding an extra filter. If it is accepted it will return an empty UP.
      * If not you will get in in return. Currently it will only be accepted by a
diff --git a/searchlib/src/vespa/searchlib/queryeval/wand/wand_parts.h b/searchlib/src/vespa/searchlib/queryeval/wand/wand_parts.h
index 9496090cca3..f0a359b98df 100644
--- a/searchlib/src/vespa/searchlib/queryeval/wand/wand_parts.h
+++ b/searchlib/src/vespa/searchlib/queryeval/wand/wand_parts.h
@@ -55,12 +55,18 @@ struct Term {
     uint32_t                 estHits;
     fef::TermFieldMatchData *matchData;
     score_t                  maxScore = 0.0; // <- only used by rise wand test
+    size_t                   old_idx = 0; // reverse-mapping used by iterator profiling
     Term(SearchIterator *s, int32_t w, uint32_t e, fef::TermFieldMatchData *tfmd) noexcept
         : search(s), weight(w), estHits(e), matchData(tfmd)
     {}
     Term() noexcept : Term(nullptr, 0, 0, nullptr){}
     Term(SearchIterator *s, int32_t w, uint32_t e) noexcept : Term(s, w, e, nullptr) {}
     Term(SearchIterator::UP s, int32_t w, uint32_t e) noexcept : Term(s.release(), w, e, nullptr) {}
+    Term copy_from(size_t idx) const noexcept {
+        Term result = *this;
+        result.old_idx = idx;
+        return result;
+    }
 };
 using Terms = std::vector<Term>;
 //-----------------------------------------------------------------------------
@@ -263,6 +269,13 @@ public:
     void unpack(uint16_t ref, uint32_t docid) { iteratorPack().unpack(ref, docid); }
     void visit_members(vespalib::ObjectVisitor &visitor) const;
     const Terms &input_terms() const { return _terms; }
+    void transform_children(auto f) {
+        iteratorPack().transform_children([&](auto itr, size_t idx){
+                                              auto ret = f(std::move(itr), _terms[idx].old_idx);
+                                              _terms[idx].search = ret.get();
+                                              return ret;
+                                          });
+    }
 };
 
 template <typename Scorer>
@@ -271,7 +284,7 @@ VectorizedIteratorTerms::VectorizedIteratorTerms(const Terms &t, const Scorer &
     : _terms()
 {
     std::vector<ref_t> order = init_state<Scorer>(TermInput(t), scorer, docIdLimit);
-    _terms = assemble([&t](ref_t ref){ return t[ref]; }, order);
+    _terms = assemble([&t](ref_t ref){ return t[ref].copy_from(ref); }, order);
     iteratorPack() = SearchIteratorPack(assemble([&t](ref_t ref){ return t[ref].search; }, order),
                                         assemble([&t](ref_t ref){ return t[ref].matchData; }, order),
                                         std::move(childrenMatchData));
diff --git a/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.cpp b/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.cpp
index 8c3e3e6be29..6f0a4b8f825 100644
--- a/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.cpp
+++ b/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.cpp
@@ -67,6 +67,9 @@ public:
     score_t get_max_score(size_t idx) const override { return _terms.maxScore(idx); }
     const Terms &getTerms() const override { return _terms.input_terms(); }
     uint32_t getN() const override { return _n; }
+    void transform_children(std::function<SearchIterator::UP(SearchIterator::UP, size_t)> f) override {
+        _terms.transform_children(std::move(f));
+    }
     void doSeek(uint32_t docid) override {
         updateThreshold(_matchParams.scores.getMinScore());
         if (IS_STRICT) {
diff --git a/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.h b/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.h
index 74c5a4bf48e..9998249cac8 100644
--- a/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.h
+++ b/searchlib/src/vespa/searchlib/queryeval/wand/weak_and_search.h
@@ -4,6 +4,7 @@
 
 #include "wand_parts.h"
 #include <vespa/searchlib/queryeval/searchiterator.h>
+#include <functional>
 
 namespace search::queryeval {
 
@@ -15,6 +16,8 @@ struct WeakAndSearch : SearchIterator {
     virtual wand::score_t get_max_score(size_t idx) const = 0;
     virtual const Terms &getTerms() const = 0;
     virtual uint32_t getN() const = 0;
+    WeakAndSearch *as_weak_and() noexcept override { return this; }
+    virtual void transform_children(std::function<SearchIterator::UP(SearchIterator::UP, size_t)> f) = 0;
     void visitMembers(vespalib::ObjectVisitor &visitor) const override;
     template<typename Scorer>
     static SearchIterator::UP createArrayWand(const Terms &terms, const MatchParams & matchParams,