Merge pull request #8824 from vespa-engine/balder/avoid-adding-another-layer-when-injecting-geo-search

Balder/avoid adding another layer when injecting geo search

2 files changed, 94 insertions, 20 deletions

diff --git a/searchcore/src/tests/proton/matching/query_test.cpp b/searchcore/src/tests/proton/matching/query_test.cpp
index 8ecf6fd4a43..03f20876adc 100644
--- a/searchcore/src/tests/proton/matching/query_test.cpp
+++ b/searchcore/src/tests/proton/matching/query_test.cpp
@@ -56,6 +56,7 @@ using search::queryeval::Blueprint;
 using search::queryeval::FakeResult;
 using search::queryeval::FakeSearchable;
 using search::queryeval::FakeRequestContext;
+using search::queryeval::FakeBlueprint;
 using search::queryeval::FieldSpec;
 using search::queryeval::FieldSpecList;
 using search::queryeval::Searchable;
@@ -105,6 +106,7 @@ class Test : public vespalib::TestApp {
     void requireThatUnknownFieldActsEmpty();
     void requireThatIllegalFieldsAreIgnored();
     void requireThatQueryGluesEverythingTogether();
+    void requireThatLocationIsAddedTheCorrectPlace();
     void requireThatQueryAddsLocation();
     void requireThatQueryAddsLocationCutoff();
     void requireThatFakeFieldSearchDumpsDiffer();
@@ -119,7 +121,7 @@ class Test : public vespalib::TestApp {
     void requireThatSameElementIteratorsCanBeBuilt();
 
 public:
-    ~Test();
+    ~Test() override;
     int Main() override;
 };
 
@@ -129,16 +131,17 @@ public:
     TEST_DO(tearDown())
 
 void Test::setUp() {
-    _match_data.reset(0);
-    _blueprint.reset(0);
+    _match_data.reset();
+    _blueprint.reset();
 }
 
 void Test::tearDown() {
-    _match_data.reset(0);
-    _blueprint.reset(0);
+    _match_data.reset();
+    _blueprint.reset();
 }
 
 const string field = "field";
+const string loc_field = "location";
 const string resolved_field1 = "resolved1";
 const string resolved_field2 = "resolved2";
 const string unknown_field = "unknown_field";
@@ -173,6 +176,10 @@ void setupIndexEnvironments()
 
     FieldInfo attr_info(FieldType::ATTRIBUTE, CollectionType::SINGLE, field, 0);
     attribute_index_env.getFields().push_back(attr_info);
+    FieldInfo loc_field_info = FieldInfo(FieldType::ATTRIBUTE, CollectionType::SINGLE,
+                                     PositionDataType::getZCurveFieldName(loc_field), field_id + 1);
+    plain_index_env.getFields().push_back(loc_field_info);
+    attribute_index_env.getFields().push_back(loc_field_info);
 }
 
 Node::UP buildQueryTree(const ViewResolver &resolver,
@@ -689,8 +696,6 @@ void Test::requireThatQueryGluesEverythingTogether() {
 }
 
 void checkQueryAddsLocation(Test &test, const string &loc_string) {
-    const string loc_field = "location";
-
     fef_test::IndexEnvironment index_environment;
     FieldInfo field_info(FieldType::INDEX, CollectionType::SINGLE, field, 0);
     index_environment.getFields().push_back(field_info);
@@ -726,6 +731,60 @@ void checkQueryAddsLocation(Test &test, const string &loc_string) {
     }
 }
 
+template<typename T1, typename T2>
+void verifyThatRankBlueprintAndAndNotStaysOnTopAfterLocation(QueryBuilder<ProtonNodeTypes> & builder) {
+    const string loc_string = "(2,10,10,3,0,1,0,0)";
+    builder.addStringTerm("foo", field, field_id, string_weight);
+    builder.addStringTerm("bar", field, field_id, string_weight);
+    builder.addStringTerm("baz", field, field_id, string_weight);
+    std::string stackDump = StackDumpCreator::create(*builder.build());
+
+    Query query;
+    query.buildTree(stackDump, loc_field + ":" + loc_string, ViewResolver(), attribute_index_env);
+    FakeSearchContext context(42);
+    context.addIdx(0).idx(0).getFake()
+            .addResult(field, "foo", FakeResult().doc(1));
+    context.setLimit(42);
+
+    query.setWhiteListBlueprint(std::make_unique<SimpleBlueprint>(SimpleResult()));
+
+    FakeRequestContext requestContext;
+    MatchDataLayout mdl;
+    query.reserveHandles(requestContext, context, mdl);
+    const IntermediateBlueprint * root = dynamic_cast<const T1 *>(query.peekRoot());
+    ASSERT_TRUE(root != nullptr);
+    EXPECT_EQUAL(2u, root->childCnt());
+    const IntermediateBlueprint * second = dynamic_cast<const T2 *>(&root->getChild(0));
+    ASSERT_TRUE(second != nullptr);
+    EXPECT_EQUAL(2u, second->childCnt());
+    auto first = dynamic_cast<const AndBlueprint *>(&second->getChild(0));
+    ASSERT_TRUE(first != nullptr);
+    EXPECT_EQUAL(2u, first->childCnt());
+    EXPECT_TRUE(dynamic_cast<const AndBlueprint *>(&first->getChild(0)));
+    auto bottom = dynamic_cast<const AndBlueprint *>(&first->getChild(0));
+    EXPECT_EQUAL(2u, bottom->childCnt());
+    EXPECT_TRUE(dynamic_cast<const FakeBlueprint *>(&bottom->getChild(0)));
+    EXPECT_TRUE(dynamic_cast<const FakeBlueprint *>(&bottom->getChild(1)));
+    EXPECT_TRUE(dynamic_cast<const SimpleBlueprint *>(&first->getChild(1)));
+    EXPECT_TRUE(dynamic_cast<const FakeBlueprint *>(&second->getChild(1)));
+    EXPECT_TRUE(dynamic_cast<const FakeBlueprint *>(&root->getChild(1)));
+}
+
+void Test::requireThatLocationIsAddedTheCorrectPlace() {
+    {
+        QueryBuilder<ProtonNodeTypes> builder;
+        builder.addRank(2);
+        builder.addAndNot(2);
+        verifyThatRankBlueprintAndAndNotStaysOnTopAfterLocation<RankBlueprint, AndNotBlueprint>(builder);
+    }
+    {
+        QueryBuilder<ProtonNodeTypes> builder;
+        builder.addAndNot(2);
+        builder.addRank(2);
+        verifyThatRankBlueprintAndAndNotStaysOnTopAfterLocation<AndNotBlueprint, RankBlueprint>(builder);
+    }
+}
+
 void Test::requireThatQueryAddsLocation() {
     checkQueryAddsLocation(*this, "(2,10,10,3,0,1,0,0)");
 }
@@ -817,8 +876,8 @@ void Test::requireThatWeakAndBlueprintsAreCreatedCorrectly() {
     wand.accept(reserve_visitor);
 
     Blueprint::UP blueprint = BlueprintBuilder::build(requestContext, wand, context);
-    WeakAndBlueprint *wbp = dynamic_cast<WeakAndBlueprint*>(blueprint.get());
-    ASSERT_TRUE(wbp != 0);
+    auto *wbp = dynamic_cast<WeakAndBlueprint*>(blueprint.get());
+    ASSERT_TRUE(wbp != nullptr);
     ASSERT_EQUAL(2u, wbp->getWeights().size());
     ASSERT_EQUAL(2u, wbp->childCnt());
     EXPECT_EQUAL(123u, wbp->getN());
@@ -851,7 +910,7 @@ void Test::requireThatParallelWandBlueprintsAreCreatedCorrectly() {
     wand.accept(reserve_visitor);
 
     Blueprint::UP blueprint = BlueprintBuilder::build(requestContext, wand, context);
-    ParallelWeakAndBlueprint *wbp = dynamic_cast<ParallelWeakAndBlueprint*>(blueprint.get());
+    auto *wbp = dynamic_cast<ParallelWeakAndBlueprint*>(blueprint.get());
     ASSERT_TRUE(wbp != nullptr);
     EXPECT_EQUAL(9000, wbp->getScoreThreshold());
     EXPECT_EQUAL(1.25, wbp->getThresholdBoostFactor());
@@ -913,7 +972,7 @@ void verifyThatRankBlueprintAndAndNotStaysOnTopAfterWhiteListing(QueryBuilder<Pr
     const IntermediateBlueprint * second = dynamic_cast<const T2 *>(&root->getChild(0));
     ASSERT_TRUE(second != nullptr);
     EXPECT_EQUAL(2u, second->childCnt());
-    const AndBlueprint * first = dynamic_cast<const AndBlueprint *>(&second->getChild(0));
+    auto first = dynamic_cast<const AndBlueprint *>(&second->getChild(0));
     ASSERT_TRUE(first != nullptr);
     EXPECT_EQUAL(2u, first->childCnt());
     EXPECT_TRUE(dynamic_cast<const SourceBlenderBlueprint *>(&first->getChild(0)));
@@ -956,7 +1015,7 @@ void
 Test::requireThatSameElementTermsAreProperlyPrefixed()
 {
     search::query::Node::UP query = make_same_element_stack_dump("", "");
-    search::query::SameElement * root = dynamic_cast<search::query::SameElement *>(query.get());
+    auto * root = dynamic_cast<search::query::SameElement *>(query.get());
     EXPECT_EQUAL(root->getView(), "");
     EXPECT_EQUAL(root->getChildren().size(), 2u);
     EXPECT_EQUAL(dynamic_cast<ProtonStringTerm *>(root->getChildren()[0])->getView(), "f1");
@@ -1035,6 +1094,7 @@ Test::Main()
     TEST_CALL(requireThatUnknownFieldActsEmpty);
     TEST_CALL(requireThatIllegalFieldsAreIgnored);
     TEST_CALL(requireThatQueryGluesEverythingTogether);
+    TEST_CALL(requireThatLocationIsAddedTheCorrectPlace);
     TEST_CALL(requireThatQueryAddsLocation);
     TEST_CALL(requireThatQueryAddsLocationCutoff);
     TEST_CALL(requireThatFakeFieldSearchDumpsDiffer);
diff --git a/searchcore/src/vespa/searchcore/proton/matching/query.cpp b/searchcore/src/vespa/searchcore/proton/matching/query.cpp
index a8f75928a6d..351de63de37 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/query.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/query.cpp
@@ -40,7 +40,25 @@ using std::vector;
 namespace proton::matching {
 
 namespace {
-void AddLocationNode(const string &location_str, Node::UP &query_tree, Location &fef_location) {
+
+Node::UP
+inject(Node::UP query, Node::UP to_inject) {
+    if (auto * my_and = dynamic_cast<search::query::And *>(query.get())) {
+        my_and->append(std::move(to_inject));
+    } else  if (dynamic_cast<search::query::Rank *>(query.get()) || dynamic_cast<search::query::AndNot *>(query.get())) {
+        search::query::Intermediate & root = static_cast<search::query::Intermediate &>(*query);
+        root.prepend(inject(root.stealFirst(), std::move(to_inject)));
+    } else {
+        auto new_root = std::make_unique<ProtonAnd>();
+        new_root->append(std::move(query));
+        new_root->append(std::move(to_inject));
+        query = std::move(new_root);
+    }
+    return query;
+}
+
+void
+addLocationNode(const string &location_str, Node::UP &query_tree, Location &fef_location) {
     if (location_str.empty()) {
         return;
     }
@@ -61,20 +79,16 @@ void AddLocationNode(const string &location_str, Node::UP &query_tree, Location
     int32_t id = -1;
     Weight weight(100);
 
-    auto new_base = std::make_unique<ProtonAnd>();
-    new_base->append(std::move(query_tree));
-
     if (locationSpec.getRankOnDistance()) {
-        new_base->append(std::make_unique<ProtonLocationTerm>(loc, view, id, weight));
+        query_tree = inject(std::move(query_tree), std::make_unique<ProtonLocationTerm>(loc, view, id, weight));
         fef_location.setAttribute(view);
         fef_location.setXPosition(locationSpec.getX());
         fef_location.setYPosition(locationSpec.getY());
         fef_location.setXAspect(locationSpec.getXAspect());
         fef_location.setValid(true);
     } else if (locationSpec.getPruneOnDistance()) {
-        new_base->append(std::make_unique<ProtonLocationTerm>(loc, view, id, weight));
+        query_tree = inject(std::move(query_tree), std::make_unique<ProtonLocationTerm>(loc, view, id, weight));
     }
-    query_tree = std::move(new_base);
 }
 
 IntermediateBlueprint *
@@ -111,7 +125,7 @@ Query::buildTree(vespalib::stringref stack, const string &location,
     if (_query_tree) {
         SameElementModifier prefixSameElementSubIndexes;
         _query_tree->accept(prefixSameElementSubIndexes);
-        AddLocationNode(location, _query_tree, _location);
+        addLocationNode(location, _query_tree, _location);
         ResolveViewVisitor resolve_visitor(resolver, indexEnv);
         _query_tree->accept(resolve_visitor);
         return true;