update MultiPhrasePrefixQuery to match phrase terms

Signed-off-by: Nicholas Walter Knize <nknize@apache.org>
opensearch-project · Jan 30, 2023 · 856560e · 856560e
1 parent 89f64e9
commit 856560e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/...is-common/src/test/java/org/opensearch/analysis/common/HighlighterWithAnalyzersTests.java b/...is-common/src/test/java/org/opensearch/analysis/common/HighlighterWithAnalyzersTests.java
@@ -50,6 +50,7 @@
 import static org.opensearch.client.Requests.searchRequest;
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.opensearch.index.query.QueryBuilders.boolQuery;
+import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
 import static org.opensearch.index.query.QueryBuilders.matchPhrasePrefixQuery;
 import static org.opensearch.index.query.QueryBuilders.matchPhraseQuery;
 import static org.opensearch.index.query.QueryBuilders.matchQuery;
@@ -61,6 +62,7 @@
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.startsWith;
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
 
 public class HighlighterWithAnalyzersTests extends OpenSearchIntegTestCase {
     @Override
@@ -270,9 +272,13 @@ public void testPhrasePrefix() throws IOException {
         refresh();
         logger.info("--> highlighting and searching on field0");
 
-        SearchSourceBuilder source = searchSource().query(matchPhrasePrefixQuery("field0", "bro"))
-            .highlighter(highlight().field("field0").order("score").preTags("<x>").postTags("</x>"));
+        SearchSourceBuilder source = searchSource().query(matchAllQuery());
         SearchResponse searchResponse = client().search(searchRequest("first_test_index").source(source)).actionGet();
+        assertHitCount(searchResponse, 2);
+
+        source = searchSource().query(matchPhrasePrefixQuery("field0", "bro"))
+            .highlighter(highlight().field("field0").order("score").preTags("<x>").postTags("</x>"));
+        searchResponse = client().search(searchRequest("first_test_index").source(source)).actionGet();
 
         assertHighlight(searchResponse, 0, "field0", 0, 1, equalTo("The quick <x>brown</x> fox jumps over the lazy dog"));
 
@@ -415,6 +421,7 @@ public void testPhrasePrefix() throws IOException {
     public static XContentBuilder type1TermVectorMapping() throws IOException {
         return XContentFactory.jsonBuilder()
             .startObject()
+            .startObject("_doc")
             .startObject("properties")
             .startObject("field1")
             .field("type", "text")
@@ -425,6 +432,7 @@ public static XContentBuilder type1TermVectorMapping() throws IOException {
             .field("term_vector", "with_positions_offsets")
             .endObject()
             .endObject()
+            .endObject()
             .endObject();
     }
 }
diff --git a/server/src/main/java/org/opensearch/common/lucene/search/MultiPhrasePrefixQuery.java b/server/src/main/java/org/opensearch/common/lucene/search/MultiPhrasePrefixQuery.java
@@ -325,6 +325,20 @@ private boolean termArraysEquals(List<Term[]> termArrays1, List<Term[]> termArra
 
     @Override
     public void visit(QueryVisitor visitor) {
-        visitor.visitLeaf(this);
+        if (visitor.acceptField(field)) {
+            // set by the subvisitor because the default returns, *this*, unless MUST_NOT is used which returns
+            // an ideal EMPTY_VISITOR
+            visitor = visitor.getSubVisitor(BooleanClause.Occur.MUST, this);
+            for (int i = 0; i < termArrays.size() - 1; ++i) {
+                if (termArrays.get(i).length == 1) {
+                    visitor.consumeTerms(this, termArrays.get(i)[0]); // use a must if we only have a single phrase
+                } else {
+                    // if more than one sub-phrase
+                    // we should match at least one
+                    QueryVisitor subVisitor = visitor.getSubVisitor(BooleanClause.Occur.SHOULD, this);
+                    subVisitor.consumeTerms(this, termArrays.get(i));
+                }
+            }
+        }
     }
 }