Adding indexordocvalue to custom query scorer

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Oct 13, 2023 · 29f79fe · 29f79fe
1 parent 292429b
commit 29f79fe
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/search/fetch/subphase/highlight/CustomQueryScorer.java b/server/src/main/java/org/opensearch/search/fetch/subphase/highlight/CustomQueryScorer.java
@@ -33,6 +33,7 @@
 package org.opensearch.search.fetch.subphase.highlight;
 
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.highlight.QueryScorer;
 import org.apache.lucene.search.highlight.WeightedSpanTerm;
@@ -104,6 +105,8 @@ protected void extract(Query query, float boost, Map<String, WeightedSpanTerm> t
                 super.extract(((FunctionScoreQuery) query).getSubQuery(), boost, terms);
             } else if (query instanceof OpenSearchToParentBlockJoinQuery) {
                 super.extract(((OpenSearchToParentBlockJoinQuery) query).getChildQuery(), boost, terms);
+            } else if (query instanceof IndexOrDocValuesQuery) {
+                super.extract(((IndexOrDocValuesQuery) query).getIndexQuery(), boost, terms);
             } else {
                 super.extract(query, boost, terms);
             }

diff --git a/server/src/test/java/org/opensearch/index/query/MultiMatchQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/MultiMatchQueryBuilderTests.java
@@ -304,10 +304,16 @@ public void testToQueryBooleanPrefixMultipleFields() throws IOException {
                 } else if (disjunct instanceof PrefixQuery) {
                     final PrefixQuery secondDisjunct = (PrefixQuery) disjunct;
                     assertThat(secondDisjunct.getPrefix(), equalTo(new Term(KEYWORD_FIELD_NAME, "foo bar")));
+                } else if (disjunct instanceof IndexOrDocValuesQuery) {
+                    final IndexOrDocValuesQuery iodvqDisjunct = (IndexOrDocValuesQuery) disjunct;
+                    assertThat(iodvqDisjunct.getIndexQuery().toString(), equalTo("mapped_string_2:foo bar*"));
                 } else {
                     throw new AssertionError();
                 }
-                assertThat(disjunct, either(instanceOf(BooleanQuery.class)).or(instanceOf(PrefixQuery.class)));
+                assertThat(
+                    disjunct,
+                    either(instanceOf(BooleanQuery.class)).or(instanceOf(PrefixQuery.class)).or(instanceOf(IndexOrDocValuesQuery.class))
+                );
             }
         }
     }