Skip to content

Commit f44b0d9

Browse files
author
girish jeyakumar
committed
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 Signed-off-by: Girish Jeyakumar <girishjeyakumar@example.com>
1 parent f4ca647 commit f44b0d9

File tree

5 files changed

+179
-7
lines changed

5 files changed

+179
-7
lines changed

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,3 +101,93 @@ setup:
101101
- match: { hits.hits.0.matched_queries.match_field_2: 10 }
102102
- length: { hits.hits.1.matched_queries: 1 }
103103
- match: { hits.hits.1.matched_queries.match_field_1: 1 }
104+
105+
---
106+
107+
"named queries in rescore":
108+
- do:
109+
indices.create:
110+
index: test
111+
112+
- do:
113+
bulk:
114+
refresh: true
115+
body:
116+
- '{ "index" : { "_index" : "test_1", "_id" : "1" } }'
117+
- '{"field" : 1, "title": "hello world" }'
118+
- '{ "index" : { "_index" : "test_1", "_id" : "2" } }'
119+
- '{"field" : 2, "title": "hello universe" }'
120+
121+
- do:
122+
search:
123+
index: test_1
124+
body:
125+
query:
126+
match: {
127+
field: {
128+
query: 1,
129+
_name: main_query
130+
}
131+
}
132+
rescore:
133+
window_size: 10
134+
query:
135+
rescore_query:
136+
match: {
137+
title: {
138+
query: "hello",
139+
_name: rescore_query
140+
}
141+
}
142+
query_weight: 0.5
143+
rescore_query_weight: 1.5
144+
145+
- match: { hits.total.value: 1 }
146+
- length: { hits.hits.0.matched_queries: 2 }
147+
- match: { hits.hits.0.matched_queries: [ "main_query", "rescore_query" ] }
148+
149+
---
150+
151+
"named queries in rescore with scores":
152+
- do:
153+
indices.create:
154+
index: test
155+
156+
- do:
157+
bulk:
158+
refresh: true
159+
body:
160+
- '{ "index" : { "_index" : "test_1", "_id" : "1" } }'
161+
- '{"field" : 1, "title": "hello world" }'
162+
- '{ "index" : { "_index" : "test_1", "_id" : "2" } }'
163+
- '{"field" : 2, "title": "hello universe" }'
164+
165+
- do:
166+
search:
167+
include_named_queries_score: true
168+
index: test_1
169+
body:
170+
query:
171+
match: {
172+
field: {
173+
query: 1,
174+
_name: main_query
175+
}
176+
}
177+
rescore:
178+
window_size: 10
179+
query:
180+
rescore_query:
181+
match: {
182+
title: {
183+
query: "hello",
184+
_name: rescore_query
185+
}
186+
}
187+
query_weight: 0.5
188+
rescore_query_weight: 1.5
189+
190+
- match: { hits.total.value: 1 }
191+
- length: { hits.hits.0.matched_queries: 2 }
192+
- gte: { hits.hits.0.matched_queries.main_query: 0.0 }
193+
- gte: { hits.hits.0.matched_queries.rescore_query: 0.0 }

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@
4141
import org.opensearch.search.fetch.FetchContext;
4242
import org.opensearch.search.fetch.FetchSubPhase;
4343
import org.opensearch.search.fetch.FetchSubPhaseProcessor;
44+
import org.opensearch.search.rescore.RescoreContext;
45+
import org.opensearch.search.rescore.QueryRescorer;
4446

4547
import java.io.IOException;
4648
import java.util.ArrayList;
@@ -65,6 +67,22 @@ public FetchSubPhaseProcessor getProcessor(FetchContext context) throws IOExcept
6567
if (context.parsedPostFilter() != null) {
6668
namedQueries.putAll(context.parsedPostFilter().namedFilters());
6769
}
70+
71+
// Also collect named queries from rescore contexts
72+
// This ensures that named queries inside rescore.query.rescore_query are included in matched_queries
73+
if (context.rescore() != null) {
74+
for (RescoreContext rescoreContext : context.rescore()) {
75+
// For QueryRescoreContext, access the ParsedQuery to get named queries
76+
// This is consistent with how main query and post filter are handled
77+
if (rescoreContext instanceof QueryRescorer.QueryRescoreContext) {
78+
QueryRescorer.QueryRescoreContext queryRescoreContext = (QueryRescorer.QueryRescoreContext) rescoreContext;
79+
if (queryRescoreContext.parsedQuery() != null) {
80+
namedQueries.putAll(queryRescoreContext.parsedQuery().namedFilters());
81+
}
82+
}
83+
}
84+
}
85+
6886
if (namedQueries.isEmpty()) {
6987
return null;
7088
}

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@
3838
import org.apache.lucene.search.ScoreDoc;
3939
import org.apache.lucene.search.TopDocs;
4040

41+
import org.opensearch.search.rescore.QueryRescorer.QueryRescoreContext;
42+
4143
import java.io.IOException;
4244
import java.util.Arrays;
4345
import java.util.Collections;
@@ -190,7 +192,7 @@ private TopDocs combine(TopDocs in, TopDocs resorted, QueryRescoreContext ctx) {
190192
* @opensearch.internal
191193
*/
192194
public static class QueryRescoreContext extends RescoreContext {
193-
private Query query;
195+
private org.opensearch.index.query.ParsedQuery parsedQuery;
194196
private float queryWeight = 1.0f;
195197
private float rescoreQueryWeight = 1.0f;
196198
private QueryRescoreMode scoreMode;
@@ -200,17 +202,21 @@ public QueryRescoreContext(int windowSize) {
200202
this.scoreMode = QueryRescoreMode.Total;
201203
}
202204

203-
public void setQuery(Query query) {
204-
this.query = query;
205+
public void setParsedQuery(org.opensearch.index.query.ParsedQuery parsedQuery) {
206+
this.parsedQuery = parsedQuery;
207+
}
208+
209+
public org.opensearch.index.query.ParsedQuery parsedQuery() {
210+
return parsedQuery;
205211
}
206212

207213
@Override
208214
public List<Query> getQueries() {
209-
return Collections.singletonList(query);
215+
return Collections.singletonList(query());
210216
}
211217

212218
public Query query() {
213-
return query;
219+
return parsedQuery != null ? parsedQuery.query() : null;
214220
}
215221

216222
public float queryWeight() {

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,9 @@
4242
import org.opensearch.index.query.QueryBuilder;
4343
import org.opensearch.index.query.QueryRewriteContext;
4444
import org.opensearch.index.query.QueryShardContext;
45+
import org.opensearch.index.query.ParsedQuery;
4546
import org.opensearch.search.rescore.QueryRescorer.QueryRescoreContext;
47+
import org.apache.lucene.search.Query;
4648

4749
import java.io.IOException;
4850
import java.util.Locale;
@@ -190,8 +192,12 @@ public static QueryRescorerBuilder fromXContent(XContentParser parser) throws IO
190192
@Override
191193
public QueryRescoreContext innerBuildContext(int windowSize, QueryShardContext context) throws IOException {
192194
QueryRescoreContext queryRescoreContext = new QueryRescoreContext(windowSize);
193-
// query is rewritten at this point already
194-
queryRescoreContext.setQuery(queryBuilder.toQuery(context));
195+
196+
// Build the ParsedQuery to capture both the Lucene Query and named queries
197+
// This is consistent with how main query and post filter are handled
198+
ParsedQuery parsedQuery = context.toQuery(queryBuilder);
199+
queryRescoreContext.setParsedQuery(parsedQuery);
200+
195201
queryRescoreContext.setQueryWeight(this.queryWeight);
196202
queryRescoreContext.setRescoreQueryWeight(this.rescoreQueryWeight);
197203
queryRescoreContext.setScoreMode(this.scoreMode);

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,58 @@ public void testRescoreQueryNull() throws IOException {
202202
assertEquals("rescore_query cannot be null", e.getMessage());
203203
}
204204

205+
/**
206+
* Test that named queries from rescore contexts are properly captured in ParsedQuery
207+
*/
208+
public void testRescoreNamedQueries() throws IOException {
209+
final long nowInMillis = randomNonNegativeLong();
210+
Settings indexSettings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build();
211+
IndexSettings idxSettings = IndexSettingsModule.newIndexSettings(randomAlphaOfLengthBetween(1, 10), indexSettings);
212+
213+
QueryShardContext mockShardContext = new QueryShardContext(
214+
0,
215+
idxSettings,
216+
BigArrays.NON_RECYCLING_INSTANCE,
217+
null,
218+
null,
219+
null,
220+
null,
221+
null,
222+
xContentRegistry(),
223+
namedWriteableRegistry,
224+
null,
225+
null,
226+
() -> nowInMillis,
227+
null,
228+
null,
229+
() -> true,
230+
null
231+
) {
232+
@Override
233+
public MappedFieldType fieldMapper(String name) {
234+
TextFieldMapper.Builder builder = new TextFieldMapper.Builder(name, createDefaultIndexAnalyzers());
235+
return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
236+
}
237+
};
238+
239+
// Create a query builder with a name
240+
QueryBuilder namedQueryBuilder = new MatchAllQueryBuilder().queryName("test_rescore_query");
241+
QueryRescorerBuilder rescoreBuilder = new QueryRescorerBuilder(namedQueryBuilder);
242+
243+
// Build the rescore context
244+
QueryRescoreContext rescoreContext = (QueryRescoreContext) rescoreBuilder.buildContext(mockShardContext);
245+
246+
// Verify that the ParsedQuery was captured and contains named queries
247+
assertNotNull(rescoreContext.parsedQuery());
248+
assertNotNull(rescoreContext.parsedQuery().namedFilters());
249+
assertEquals(1, rescoreContext.parsedQuery().namedFilters().size());
250+
assertTrue(rescoreContext.parsedQuery().namedFilters().containsKey("test_rescore_query"));
251+
assertNotNull(rescoreContext.parsedQuery().namedFilters().get("test_rescore_query"));
252+
253+
// Verify that the Lucene query is still accessible
254+
assertNotNull(rescoreContext.query());
255+
}
256+
205257
class AlwaysRewriteQueryBuilder extends MatchAllQueryBuilder {
206258

207259
@Override

0 commit comments

Comments
 (0)