Fix: Include named queries from rescore contexts in matched_queries array

- Modified QueryRescoreContext to store ParsedQuery instead of just Query
- Updated QueryRescorerBuilder to use context.toQuery() for capturing named queries
- Enhanced MatchedQueriesPhase to collect named queries from all rescore contexts
- Added comprehensive tests for both unit and REST API scenarios
- Resolves inconsistency with Elasticsearch behavior where rescore named queries were not surfaced
- Cleaned up unnecessary comments for better code readability

Signed-off-by: Girish Jeyakumar <[email protected]>

Author: Girish Jeyakumar

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/fetch/subphase/MatchedQueriesPhase.java

@@ -41,6 +41,8 @@
 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 org.opensearch.search.rescore.QueryRescorer;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -65,6 +67,17 @@ 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 instanceof QueryRescorer.QueryRescoreContext) {
+                    QueryRescorer.QueryRescoreContext queryRescoreContext = (QueryRescorer.QueryRescoreContext) rescoreContext;
+                    if (queryRescoreContext.parsedQuery() != null) {
+                        namedQueries.putAll(queryRescoreContext.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;
@@ -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,21 @@ 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;
+        }
+
+        public org.opensearch.index.query.ParsedQuery parsedQuery() {
+            return parsedQuery;
         }
 
         @Override
         public List<Query> getQueries() {
-            return Collections.singletonList(query);
+            return Collections.singletonList(query());
         }
 
         public Query query() {
-            return query;
+            return parsedQuery != null ? parsedQuery.query() : null;
         }
 
         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,10 @@ 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));
+
+        ParsedQuery parsedQuery = context.toQuery(queryBuilder);
+        queryRescoreContext.setParsedQuery(parsedQuery);
+
         queryRescoreContext.setQueryWeight(this.queryWeight);
         queryRescoreContext.setRescoreQueryWeight(this.rescoreQueryWeight);
         queryRescoreContext.setScoreMode(this.scoreMode);

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

@@ -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