Include named queries from rescore contexts in matched_queries array

CHANGELOG.md

@@ -24,6 +24,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Extend Approximation Framework to other numeric types ([#18530](https://github.com/opensearch-project/OpenSearch/issues/18530))
 - Add Semantic Version field type mapper and extensive unit tests([#18454](https://github.com/opensearch-project/OpenSearch/pull/18454))
 - Pass index settings to system ingest processor factories. ([#18708](https://github.com/opensearch-project/OpenSearch/pull/18708))
+- Include named queries from rescore contexts in matched_queries array ([#18697](https://github.com/opensearch-project/OpenSearch/pull/18697))
 
 ### Changed
 - Update Subject interface to use CheckedRunnable ([#18570](https://github.com/opensearch-project/OpenSearch/issues/18570))

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

@@ -101,3 +101,101 @@ 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":
+  - skip:
+      version: " - 3.1.99"
+      reason: "named queries in rescore is supported in 3.2.0 and above"
+
+  - 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":
+  - skip:
+      version: " - 3.1.99"
+      reason: "named queries in rescore is supported in 3.2.0 and above"
+
+  - 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

@@ -35,10 +35,10 @@
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.CollectionStatistics;
 import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.TermStatistics;
 import org.opensearch.core.tasks.TaskCancelledException;
+import org.opensearch.index.query.ParsedQuery;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.rescore.RescoreContext;
 
@@ -86,8 +86,8 @@
 
             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 (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,12 @@
         if (context.parsedPostFilter() != null) {
             namedQueries.putAll(context.parsedPostFilter().namedFilters());
         }
+        if (context.rescore() != null) {
+            for (RescoreContext rescoreContext : context.rescore()) {
+                rescoreContext.getParsedQueries().forEach(query -> namedQueries.putAll(query.namedFilters()));
+            }
+        }
+
         if (namedQueries.isEmpty()) {
             return null;
         }

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

@@ -37,6 +37,8 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.TopDocs;
+import org.opensearch.index.query.ParsedQuery;
+import org.opensearch.search.rescore.QueryRescorer.QueryRescoreContext;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -67,7 +69,7 @@
 
         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 @@
             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 @@
      * @opensearch.internal
      */
     public static class QueryRescoreContext extends RescoreContext {
-        private Query query;
+        private ParsedQuery parsedQuery;
         private float queryWeight = 1.0f;
         private float rescoreQueryWeight = 1.0f;
         private QueryRescoreMode scoreMode;
@@ -200,17 +202,22 @@
             this.scoreMode = QueryRescoreMode.Total;
         }
 
-        public void setQuery(Query query) {
-            this.query = query;
+        public void setParsedQuery(ParsedQuery parsedQuery) {
+            this.parsedQuery = parsedQuery;
+        }
+
+        public ParsedQuery parsedQuery() {
+            return parsedQuery;
         }
 
         @Override
         public List<Query> getQueries() {
-            return Collections.singletonList(query);
+            return parsedQuery != null ? Collections.singletonList(parsedQuery.query()) : Collections.emptyList();
         }
 
-        public Query query() {
-            return query;
+        @Override
+        public List<ParsedQuery> getParsedQueries() {
+            return parsedQuery != null ? Collections.singletonList(parsedQuery) : Collections.emptyList();
         }
 
         public float queryWeight() {

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

@@ -190,8 +190,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;
@@ -93,4 +94,11 @@
     public List<Query> getQueries() {
         return Collections.emptyList();
     }
+
+    /**
+     * 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