Include named queries from rescore contexts in matched_queries array

rest-api-spec/src/main/resources/rest-api-spec/test/search/350_matched_queries.yml

@@ -101,3 +101,93 @@ setup:
   - match: { hits.hits.0.matched_queries.match_field_2: 10 }
   - length: { hits.hits.1.matched_queries: 1 }
   - match: { hits.hits.1.matched_queries.match_field_1: 1 }
+
+---
+
+"named queries in rescore":
+  - do:
+      indices.create:
+        index: test
+
+  - do:
+      bulk:
+        refresh: true
+        body:
+          - '{ "index" : { "_index" : "test_1", "_id" : "1" } }'
+          - '{"field" : 1, "title": "hello world" }'
+          - '{ "index" : { "_index" : "test_1", "_id" : "2" } }'
+          - '{"field" : 2, "title": "hello universe" }'
+
+  - do:
+      search:
+        index: test_1
+        body:
+          query:
+            match: {
+              field: {
+                query: 1,
+                _name: main_query
+              }
+            }
+          rescore:
+            window_size: 10
+            query:
+              rescore_query:
+                match: {
+                  title: {
+                    query: "hello",
+                    _name: rescore_query
+                  }
+                }
+              query_weight: 0.5
+              rescore_query_weight: 1.5
+
+  - match: { hits.total.value: 1 }
+  - length: { hits.hits.0.matched_queries: 2 }
+  - match: { hits.hits.0.matched_queries: [ "main_query", "rescore_query" ] }
+
+---
+
+"named queries in rescore with scores":
+  - do:
+      indices.create:
+        index: test
+
+  - do:
+      bulk:
+        refresh: true
+        body:
+          - '{ "index" : { "_index" : "test_1", "_id" : "1" } }'
+          - '{"field" : 1, "title": "hello world" }'
+          - '{ "index" : { "_index" : "test_1", "_id" : "2" } }'
+          - '{"field" : 2, "title": "hello universe" }'
+
+  - do:
+      search:
+        include_named_queries_score: true
+        index: test_1
+        body:
+          query:
+            match: {
+              field: {
+                query: 1,
+                _name: main_query
+              }
+            }
+          rescore:
+            window_size: 10
+            query:
+              rescore_query:
+                match: {
+                  title: {
+                    query: "hello",
+                    _name: rescore_query
+                  }
+                }
+              query_weight: 0.5
+              rescore_query_weight: 1.5
+
+  - match: { hits.total.value: 1 }
+  - length: { hits.hits.0.matched_queries: 2 }
+  - gte: { hits.hits.0.matched_queries.main_query: 0.0 }
+  - gte: { hits.hits.0.matched_queries.rescore_query: 0.0 }

server/src/main/java/org/opensearch/search/dfs/DfsPhase.java

@@ -86,8 +86,8 @@ public CollectionStatistics collectionStatistics(String field) throws IOExceptio
 
             searcher.createWeight(context.searcher().rewrite(context.query()), ScoreMode.COMPLETE, 1);
             for (RescoreContext rescoreContext : context.rescore()) {
-                for (Query query : rescoreContext.getQueries()) {
-                    searcher.createWeight(context.searcher().rewrite(query), ScoreMode.COMPLETE, 1);
+                for (org.opensearch.index.query.ParsedQuery parsedQuery : rescoreContext.getParsedQueries()) {
+                    searcher.createWeight(context.searcher().rewrite(parsedQuery.query()), ScoreMode.COMPLETE, 1);
                 }
             }
 

server/src/main/java/org/opensearch/search/fetch/subphase/MatchedQueriesPhase.java

@@ -41,6 +41,7 @@
 import org.opensearch.search.fetch.FetchContext;
 import org.opensearch.search.fetch.FetchSubPhase;
 import org.opensearch.search.fetch.FetchSubPhaseProcessor;
+import org.opensearch.search.rescore.RescoreContext;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -65,6 +66,14 @@ public FetchSubPhaseProcessor getProcessor(FetchContext context) throws IOExcept
         if (context.parsedPostFilter() != null) {
             namedQueries.putAll(context.parsedPostFilter().namedFilters());
         }
+        if (context.rescore() != null) {
+            for (RescoreContext rescoreContext : context.rescore()) {
+                if (rescoreContext.parsedQuery() != null) {
+                    namedQueries.putAll(rescoreContext.parsedQuery().namedFilters());
+                }
+            }
+        }
+
         if (namedQueries.isEmpty()) {
             return null;
         }

server/src/main/java/org/opensearch/search/rescore/QueryRescorer.java

@@ -38,6 +38,8 @@
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.TopDocs;
 
+import org.opensearch.search.rescore.QueryRescorer.QueryRescoreContext;
+
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
@@ -67,7 +69,7 @@ public TopDocs rescore(TopDocs topDocs, IndexSearcher searcher, RescoreContext r
 
         final QueryRescoreContext rescore = (QueryRescoreContext) rescoreContext;
 
-        org.apache.lucene.search.Rescorer rescorer = new org.apache.lucene.search.QueryRescorer(rescore.query()) {
+        org.apache.lucene.search.Rescorer rescorer = new org.apache.lucene.search.QueryRescorer(rescore.parsedQuery().query()) {
 
             @Override
             protected float combine(float firstPassScore, boolean secondPassMatches, float secondPassScore) {
@@ -120,7 +122,7 @@ public Explanation explain(int topLevelDocId, IndexSearcher searcher, RescoreCon
             prim = Explanation.noMatch("First pass did not match", sourceExplanation);
         }
         if (rescoreContext.isRescored(topLevelDocId)) {
-            Explanation rescoreExplain = searcher.explain(rescore.query(), topLevelDocId);
+            Explanation rescoreExplain = searcher.explain(rescore.parsedQuery().query(), topLevelDocId);
             // NOTE: we don't use Lucene's Rescorer.explain because we want to insert our own description with which ScoreMode was used.
             // Maybe we should add QueryRescorer.explainCombine to Lucene?
             if (rescoreExplain != null && rescoreExplain.isMatch()) {
@@ -190,7 +192,7 @@ private TopDocs combine(TopDocs in, TopDocs resorted, QueryRescoreContext ctx) {
      * @opensearch.internal
      */
     public static class QueryRescoreContext extends RescoreContext {
-        private Query query;
+        private org.opensearch.index.query.ParsedQuery parsedQuery;
         private float queryWeight = 1.0f;
         private float rescoreQueryWeight = 1.0f;
         private QueryRescoreMode scoreMode;
@@ -200,17 +202,18 @@ public QueryRescoreContext(int windowSize) {
             this.scoreMode = QueryRescoreMode.Total;
         }
 
-        public void setQuery(Query query) {
-            this.query = query;
+        public void setParsedQuery(org.opensearch.index.query.ParsedQuery parsedQuery) {
+            this.parsedQuery = parsedQuery;
         }
 
         @Override
-        public List<Query> getQueries() {
-            return Collections.singletonList(query);
+        public org.opensearch.index.query.ParsedQuery parsedQuery() {
+            return parsedQuery;
         }
 
-        public Query query() {
-            return query;
+        @Override
+        public List<org.opensearch.index.query.ParsedQuery> getParsedQueries() {
+            return parsedQuery != null ? Collections.singletonList(parsedQuery) : Collections.emptyList();
         }
 
         public float queryWeight() {

server/src/main/java/org/opensearch/search/rescore/QueryRescorerBuilder.java

@@ -42,7 +42,9 @@
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.index.query.ParsedQuery;
 import org.opensearch.search.rescore.QueryRescorer.QueryRescoreContext;
+import org.apache.lucene.search.Query;
 
 import java.io.IOException;
 import java.util.Locale;
@@ -190,8 +192,9 @@ public static QueryRescorerBuilder fromXContent(XContentParser parser) throws IO
     @Override
     public QueryRescoreContext innerBuildContext(int windowSize, QueryShardContext context) throws IOException {
         QueryRescoreContext queryRescoreContext = new QueryRescoreContext(windowSize);
-        // query is rewritten at this point already
-        queryRescoreContext.setQuery(queryBuilder.toQuery(context));
+
+        queryRescoreContext.setParsedQuery(context.toQuery(queryBuilder));
+
         queryRescoreContext.setQueryWeight(this.queryWeight);
         queryRescoreContext.setRescoreQueryWeight(this.rescoreQueryWeight);
         queryRescoreContext.setScoreMode(this.scoreMode);

server/src/main/java/org/opensearch/search/rescore/RescoreContext.java

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.search.Query;
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.index.query.ParsedQuery;
 
 import java.util.Collections;
 import java.util.List;
@@ -88,9 +89,16 @@ public Set<Integer> getRescoredDocs() {
     }
 
     /**
-     * Returns queries associated with the rescorer
+     * Returns the parsed query for this rescore context, or null if not applicable
      */
-    public List<Query> getQueries() {
+    public ParsedQuery parsedQuery() {
+        return null;
+    }
+
+    /**
+     * Returns parsed queries associated with the rescorer
+     */
+    public List<ParsedQuery> getParsedQueries() {
         return Collections.emptyList();
     }
 }

server/src/test/java/org/opensearch/search/rescore/QueryRescorerBuilderTests.java

@@ -190,7 +190,7 @@ public MappedFieldType fieldMapper(String name) {
                 : rescoreBuilder.windowSize().intValue();
             assertEquals(expectedWindowSize, rescoreContext.getWindowSize());
             Query expectedQuery = Rewriteable.rewrite(rescoreBuilder.getRescoreQuery(), mockShardContext).toQuery(mockShardContext);
-            assertEquals(expectedQuery, rescoreContext.query());
+            assertEquals(expectedQuery, rescoreContext.parsedQuery().query());
             assertEquals(rescoreBuilder.getQueryWeight(), rescoreContext.queryWeight(), Float.MIN_VALUE);
             assertEquals(rescoreBuilder.getRescoreQueryWeight(), rescoreContext.rescoreQueryWeight(), Float.MIN_VALUE);
             assertEquals(rescoreBuilder.getScoreMode(), rescoreContext.scoreMode());
@@ -202,6 +202,50 @@ public void testRescoreQueryNull() throws IOException {
         assertEquals("rescore_query cannot be null", e.getMessage());
     }
 
+    /**
+     * Test that named queries from rescore contexts are captured
+     */
+    public void testRescoreNamedQueries() throws IOException {
+        final long nowInMillis = randomNonNegativeLong();
+        Settings indexSettings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build();
+        IndexSettings idxSettings = IndexSettingsModule.newIndexSettings(randomAlphaOfLengthBetween(1, 10), indexSettings);
+
+        QueryShardContext mockShardContext = new QueryShardContext(
+            0,
+            idxSettings,
+            BigArrays.NON_RECYCLING_INSTANCE,
+            null,
+            null,
+            null,
+            null,
+            null,
+            xContentRegistry(),
+            namedWriteableRegistry,
+            null,
+            null,
+            () -> nowInMillis,
+            null,
+            null,
+            () -> true,
+            null
+        ) {
+            @Override
+            public MappedFieldType fieldMapper(String name) {
+                TextFieldMapper.Builder builder = new TextFieldMapper.Builder(name, createDefaultIndexAnalyzers());
+                return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
+            }
+        };
+
+        QueryBuilder namedQueryBuilder = new MatchAllQueryBuilder().queryName("test_rescore_query");
+        QueryRescorerBuilder rescoreBuilder = new QueryRescorerBuilder(namedQueryBuilder);
+        QueryRescoreContext rescoreContext = (QueryRescoreContext) rescoreBuilder.buildContext(mockShardContext);
+        assertNotNull(rescoreContext.parsedQuery());
+        assertNotNull(rescoreContext.parsedQuery().namedFilters());
+        assertEquals(1, rescoreContext.parsedQuery().namedFilters().size());
+        assertTrue(rescoreContext.parsedQuery().namedFilters().containsKey("test_rescore_query"));
+        assertNotNull(rescoreContext.parsedQuery().namedFilters().get("test_rescore_query"));
+    }
+
     class AlwaysRewriteQueryBuilder extends MatchAllQueryBuilder {
 
         @Override