Merge pull request #10850 from vespa-engine/havardpe/collect-same-element-matches-from-query-state

collect matching elements by using SameElement blueprints

5 files changed, 53 insertions, 11 deletions

diff --git a/searchcore/src/tests/proton/matching/same_element_builder/same_element_builder_test.cpp b/searchcore/src/tests/proton/matching/same_element_builder/same_element_builder_test.cpp
index 6583f3bc00b..3eff38fbf20 100644
--- a/searchcore/src/tests/proton/matching/same_element_builder/same_element_builder_test.cpp
+++ b/searchcore/src/tests/proton/matching/same_element_builder/same_element_builder_test.cpp
@@ -90,7 +90,7 @@ struct BuilderFixture {
     FakeRequestContext req_ctx;
     FakeSearchContext  ctx;
     SameElementBuilder builder;
-    BuilderFixture() : req_ctx(), ctx(), builder(req_ctx, ctx, false) {
+    BuilderFixture() : req_ctx(), ctx(), builder(req_ctx, ctx, "foo", false) {
         ctx.attr().tag("attr");
         ctx.addIdx(0).idx(0).getFake().tag("idx");
     }
diff --git a/searchcore/src/vespa/searchcore/proton/matching/blueprintbuilder.cpp b/searchcore/src/vespa/searchcore/proton/matching/blueprintbuilder.cpp
index 380aeb82ad2..7e55c8f778c 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/blueprintbuilder.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/blueprintbuilder.cpp
@@ -102,7 +102,7 @@ private:
     }
 
     void buildSameElement(ProtonSameElement &n) {
-        SameElementBuilder builder(_requestContext, _context, n.is_expensive());
+        SameElementBuilder builder(_requestContext, _context, n.getView(), n.is_expensive());
         for (search::query::Node *node : n.getChildren()) {
             builder.add_child(*node);
         }
diff --git a/searchcore/src/vespa/searchcore/proton/matching/docsum_matcher.cpp b/searchcore/src/vespa/searchcore/proton/matching/docsum_matcher.cpp
index 73aab5b3fca..34310371755 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/docsum_matcher.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/docsum_matcher.cpp
@@ -4,15 +4,23 @@
 #include <vespa/eval/eval/tensor.h>
 #include <vespa/eval/eval/tensor_engine.h>
 #include <vespa/vespalib/objects/nbostream.h>
+#include <vespa/searchlib/queryeval/blueprint.h>
+#include <vespa/searchlib/queryeval/intermediate_blueprints.h>
+#include <vespa/searchlib/queryeval/same_element_blueprint.h>
+#include <vespa/searchlib/queryeval/same_element_search.h>
 
 #include <vespa/log/log.h>
 LOG_SETUP(".proton.matching.docsum_matcher");
 
 using search::FeatureSet;
-using search::StructFieldMapper;
 using search::MatchingElements;
-using search::fef::RankProgram;
+using search::StructFieldMapper;
 using search::fef::FeatureResolver;
+using search::fef::RankProgram;
+using search::queryeval::AndNotBlueprint;
+using search::queryeval::Blueprint;
+using search::queryeval::IntermediateBlueprint;
+using search::queryeval::SameElementBlueprint;
 using search::queryeval::SearchIterator;
 
 namespace proton::matching {
@@ -74,6 +82,36 @@ FeatureSet::UP get_feature_set(const MatchToolsFactory &mtf,
     return retval;
 }
 
+template<typename T>
+const T *as(const Blueprint &bp) { return dynamic_cast<const T *>(&bp); }
+
+void find_matching_elements(const std::vector<uint32_t> &docs, const SameElementBlueprint &same_element, MatchingElements &result) {
+    auto search = same_element.create_same_element_search(false);
+    search->initRange(docs.front(), docs.back()+1);
+    std::vector<uint32_t> matches;
+    for (uint32_t i = 0; i < docs.size(); ++i) {
+        search->find_matching_elements(docs[i], matches);
+        if (!matches.empty()) {
+            result.add_matching_elements(docs[i], same_element.struct_field_name(), matches);
+            matches.clear();
+        }
+    }
+}
+
+void find_matching_elements(const StructFieldMapper &mapper, const std::vector<uint32_t> &docs, const Blueprint &bp, MatchingElements &result) {
+    if (auto same_element = as<SameElementBlueprint>(bp)) {
+        if (mapper.is_struct_field(same_element->struct_field_name())) {
+            find_matching_elements(docs, *same_element, result);
+        }
+    } else if (auto and_not = as<AndNotBlueprint>(bp)) {
+        find_matching_elements(mapper, docs, and_not->getChild(0), result);
+    } else if (auto intermediate = as<IntermediateBlueprint>(bp)) {
+        for (size_t i = 0; i < intermediate->childCnt(); ++i) {
+            find_matching_elements(mapper, docs, intermediate->getChild(i), result);
+        }
+    }
+}
+
 }
 
 DocsumMatcher::DocsumMatcher()
@@ -127,11 +165,13 @@ DocsumMatcher::get_rank_features() const
 MatchingElements::UP
 DocsumMatcher::get_matching_elements(const StructFieldMapper &field_mapper) const
 {
-    if (!_mtf) {
-        return std::make_unique<MatchingElements>();
+    auto result = std::make_unique<MatchingElements>();
+    if (_mtf && !field_mapper.empty()) {
+        if (const Blueprint *root = _mtf->query().peekRoot()) {
+            find_matching_elements(field_mapper, _docs, *root, *result);
+        }
     }
-    (void) field_mapper;
-    return std::make_unique<MatchingElements>();
+    return result;
 }
 
 }
diff --git a/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.cpp b/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.cpp
index 8e012c52b5f..241ab53874f 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.cpp
@@ -71,10 +71,11 @@ public:
 
 } // namespace proton::matching::<unnamed>
 
-SameElementBuilder::SameElementBuilder(const search::queryeval::IRequestContext &requestContext, ISearchContext &context, bool expensive)
+SameElementBuilder::SameElementBuilder(const search::queryeval::IRequestContext &requestContext, ISearchContext &context,
+                                       const vespalib::string &struct_field_name, bool expensive)
     : _requestContext(requestContext),
       _context(context),
-      _result(std::make_unique<SameElementBlueprint>(expensive))
+      _result(std::make_unique<SameElementBlueprint>(struct_field_name, expensive))
 {
 }
 
diff --git a/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.h b/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.h
index b9cfff8a9c0..c68bdfced99 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.h
+++ b/searchcore/src/vespa/searchcore/proton/matching/same_element_builder.h
@@ -16,7 +16,8 @@ private:
     ISearchContext                                          &_context;
     std::unique_ptr<search::queryeval::SameElementBlueprint> _result;
 public:
-    SameElementBuilder(const search::queryeval::IRequestContext &requestContext, ISearchContext &context, bool expensive);
+    SameElementBuilder(const search::queryeval::IRequestContext &requestContext, ISearchContext &context,
+                       const vespalib::string &struct_field_name, bool expensive);
     void add_child(search::query::Node &node);
     search::queryeval::Blueprint::UP build();
 };