Skip to content

Commit 0c3a313

Browse files
harshavamsiharshavamsi
authored
Fix derived query rewrite (#19496)
Signed-off-by: Harsha Vamsi Kalluri <harshavamsi096@gmail.com>
1 parent 9923617 commit 0c3a313

File tree

4 files changed

+244
-17
lines changed

4 files changed

+244
-17
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
107107
- Setting number of sharedArenaMaxPermits to 1 ([#19503](https://github.com/opensearch-project/OpenSearch/pull/19503))
108108
- Handle negative search request nodes stats ([#19340](https://github.com/opensearch-project/OpenSearch/pull/19340))
109109
- Remove unnecessary iteration per-shard in request cache cleanup ([#19263](https://github.com/opensearch-project/OpenSearch/pull/19263))
110+
- Fix derived field rewrite to handle range queries ([#19496](https://github.com/opensearch-project/OpenSearch/pull/19496))
110111

111112
### Dependencies
112113
- Bump `com.gradleup.shadow:shadow-gradle-plugin` from 8.3.5 to 8.3.9 ([#19400](https://github.com/opensearch-project/OpenSearch/pull/19400))

‎server/src/main/java/org/opensearch/index/query/DerivedFieldQuery.java‎

Lines changed: 43 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515
import org.apache.lucene.search.ConstantScoreScorer;
1616
import org.apache.lucene.search.ConstantScoreWeight;
1717
import org.apache.lucene.search.DocIdSetIterator;
18+
import org.apache.lucene.search.IndexOrDocValuesQuery;
1819
import org.apache.lucene.search.IndexSearcher;
20+
import org.apache.lucene.search.PointRangeQuery;
1921
import org.apache.lucene.search.Query;
2022
import org.apache.lucene.search.QueryVisitor;
2123
import org.apache.lucene.search.ScoreMode;
@@ -24,6 +26,7 @@
2426
import org.apache.lucene.search.TwoPhaseIterator;
2527
import org.apache.lucene.search.Weight;
2628
import org.opensearch.index.mapper.DerivedFieldValueFetcher;
29+
import org.opensearch.search.approximate.ApproximateScoreQuery;
2730
import org.opensearch.search.lookup.LeafSearchLookup;
2831
import org.opensearch.search.lookup.SearchLookup;
2932

@@ -72,22 +75,46 @@ public void visit(QueryVisitor visitor) {
7275
query.visit(visitor);
7376
}
7477

75-
// @Override
76-
// public Query rewrite(IndexSearcher indexSearcher) throws IOException {
77-
// Query rewritten = query.rewrite(indexSearcher);
78-
// if (rewritten == query) {
79-
// return this;
80-
// }
81-
// ;
82-
// return new DerivedFieldQuery(
83-
// rewritten,
84-
// valueFetcherSupplier,
85-
// searchLookup,
86-
// indexAnalyzer,
87-
// indexableFieldGenerator,
88-
// ignoreMalformed
89-
// );
90-
// }
78+
@Override
79+
public Query rewrite(IndexSearcher indexSearcher) throws IOException {
80+
if (!needsRewrite()) {
81+
return this;
82+
}
83+
Query rewritten = query.rewrite(indexSearcher);
84+
if (rewritten == query) {
85+
return this;
86+
}
87+
;
88+
return new DerivedFieldQuery(
89+
rewritten,
90+
valueFetcherSupplier,
91+
searchLookup,
92+
indexAnalyzer,
93+
indexableFieldGenerator,
94+
ignoreMalformed
95+
);
96+
}
97+
98+
private boolean needsRewrite() {
99+
if (query instanceof PointRangeQuery) {
100+
return false;
101+
}
102+
103+
if (query instanceof ApproximateScoreQuery) {
104+
Query originalQuery = ((ApproximateScoreQuery) query).getOriginalQuery();
105+
if (originalQuery instanceof IndexOrDocValuesQuery) {
106+
Query indexQuery = ((IndexOrDocValuesQuery) originalQuery).getIndexQuery();
107+
return !(indexQuery instanceof PointRangeQuery);
108+
}
109+
}
110+
111+
if (query instanceof IndexOrDocValuesQuery) {
112+
Query indexQuery = ((IndexOrDocValuesQuery) query).getIndexQuery();
113+
return !(indexQuery instanceof PointRangeQuery);
114+
}
115+
116+
return true;
117+
}
91118

92119
@Override
93120
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {

‎server/src/test/java/org/opensearch/index/mapper/DerivedFieldMapperQueryTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -551,7 +551,7 @@ public void execute() {
551551

552552
query = new MatchPhrasePrefixQueryBuilder("object_field.text_field", "document number").toQuery(queryShardContext);
553553
topDocs = searcher.search(query, 10);
554-
assertEquals(10, topDocs.totalHits.value());
554+
assertEquals(0, topDocs.totalHits.value());
555555

556556
// Multi Phrase Query
557557
query = QueryBuilders.multiMatchQuery("GET", "object_field.nested_field.sub_field_1", "object_field.keyword_field")

‎server/src/test/java/org/opensearch/index/query/DerivedFieldQueryTests.java‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
import org.apache.lucene.document.Document;
1212
import org.apache.lucene.document.Field;
13+
import org.apache.lucene.document.IntPoint;
1314
import org.apache.lucene.document.KeywordField;
1415
import org.apache.lucene.document.TextField;
1516
import org.apache.lucene.index.DirectoryReader;
@@ -18,7 +19,10 @@
1819
import org.apache.lucene.index.IndexWriterConfig;
1920
import org.apache.lucene.index.IndexableField;
2021
import org.apache.lucene.index.Term;
22+
import org.apache.lucene.search.IndexOrDocValuesQuery;
2123
import org.apache.lucene.search.IndexSearcher;
24+
import org.apache.lucene.search.MatchNoDocsQuery;
25+
import org.apache.lucene.search.Query;
2226
import org.apache.lucene.search.TermQuery;
2327
import org.apache.lucene.search.TopDocs;
2428
import org.apache.lucene.store.Directory;
@@ -181,4 +185,199 @@ public void execute() {
181185
}
182186
}
183187
}
188+
189+
public void testNeedsRewriteWithPointRangeQuery() throws IOException {
190+
// Create lucene documents
191+
List<Document> docs = new ArrayList<>();
192+
for (String[] request : raw_requests) {
193+
Document document = new Document();
194+
document.add(new TextField("raw_request", request[0], Field.Store.YES));
195+
document.add(new KeywordField("status", request[1], Field.Store.YES));
196+
docs.add(document);
197+
}
198+
199+
// Mock SearchLookup
200+
SearchLookup searchLookup = mock(SearchLookup.class);
201+
SourceLookup sourceLookup = new SourceLookup();
202+
LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
203+
when(leafLookup.source()).thenReturn(sourceLookup);
204+
205+
// Mock DerivedFieldScript.Factory
206+
DerivedFieldScript.Factory factory = (params, lookup) -> (DerivedFieldScript.LeafFactory) ctx -> {
207+
when(searchLookup.getLeafSearchLookup(ctx)).thenReturn(leafLookup);
208+
return new DerivedFieldScript(params, lookup, ctx) {
209+
@Override
210+
public void execute() {
211+
addEmittedValue(raw_requests[sourceLookup.docId()][2]);
212+
}
213+
};
214+
};
215+
216+
// Create ValueFetcher from mocked DerivedFieldScript.Factory
217+
DerivedFieldScript.LeafFactory leafFactory = factory.newFactory((new Script("")).getParams(), searchLookup);
218+
Function<Object, IndexableField> indexableFieldFunction = DerivedFieldSupportedTypes.getIndexableFieldGeneratorType(
219+
"keyword",
220+
"ip_from_raw_request"
221+
);
222+
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(leafFactory, null);
223+
224+
// Create a real PointRangeQuery using IntPoint
225+
Query pointRangeQuery = IntPoint.newRangeQuery("test_field", 0, 100);
226+
227+
// Create DerivedFieldQuery with PointRangeQuery
228+
DerivedFieldQuery derivedFieldQuery = new DerivedFieldQuery(
229+
pointRangeQuery,
230+
() -> valueFetcher,
231+
searchLookup,
232+
Lucene.STANDARD_ANALYZER,
233+
indexableFieldFunction,
234+
true
235+
);
236+
237+
// Index and Search
238+
try (Directory dir = newDirectory()) {
239+
IndexWriter iw = new IndexWriter(dir, new IndexWriterConfig(Lucene.STANDARD_ANALYZER));
240+
for (Document d : docs) {
241+
iw.addDocument(d);
242+
}
243+
try (IndexReader reader = DirectoryReader.open(iw)) {
244+
iw.close();
245+
IndexSearcher searcher = new IndexSearcher(reader);
246+
247+
// Rewrite should return the same query without attempting rewrite
248+
Query rewritten = derivedFieldQuery.rewrite(searcher);
249+
assertSame(derivedFieldQuery, rewritten);
250+
}
251+
}
252+
}
253+
254+
public void testNeedsRewriteWithIndexOrDocValuesQueryAndPointRange() throws IOException {
255+
// Create lucene documents
256+
List<Document> docs = new ArrayList<>();
257+
for (String[] request : raw_requests) {
258+
Document document = new Document();
259+
document.add(new TextField("raw_request", request[0], Field.Store.YES));
260+
document.add(new KeywordField("status", request[1], Field.Store.YES));
261+
docs.add(document);
262+
}
263+
264+
// Mock SearchLookup
265+
SearchLookup searchLookup = mock(SearchLookup.class);
266+
SourceLookup sourceLookup = new SourceLookup();
267+
LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
268+
when(leafLookup.source()).thenReturn(sourceLookup);
269+
270+
// Mock DerivedFieldScript.Factory
271+
DerivedFieldScript.Factory factory = (params, lookup) -> (DerivedFieldScript.LeafFactory) ctx -> {
272+
when(searchLookup.getLeafSearchLookup(ctx)).thenReturn(leafLookup);
273+
return new DerivedFieldScript(params, lookup, ctx) {
274+
@Override
275+
public void execute() {
276+
addEmittedValue(raw_requests[sourceLookup.docId()][2]);
277+
}
278+
};
279+
};
280+
281+
// Create ValueFetcher from mocked DerivedFieldScript.Factory
282+
DerivedFieldScript.LeafFactory leafFactory = factory.newFactory((new Script("")).getParams(), searchLookup);
283+
Function<Object, IndexableField> indexableFieldFunction = DerivedFieldSupportedTypes.getIndexableFieldGeneratorType(
284+
"keyword",
285+
"ip_from_raw_request"
286+
);
287+
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(leafFactory, null);
288+
289+
// Create real queries: PointRangeQuery -> IndexOrDocValuesQuery
290+
Query pointRangeQuery = IntPoint.newRangeQuery("test_field", 0, 100);
291+
Query dvQuery = new MatchNoDocsQuery();
292+
IndexOrDocValuesQuery indexOrDocValuesQuery = new IndexOrDocValuesQuery(pointRangeQuery, dvQuery);
293+
294+
// Create DerivedFieldQuery with IndexOrDocValuesQuery with PointRangeQuery
295+
DerivedFieldQuery derivedFieldQuery = new DerivedFieldQuery(
296+
indexOrDocValuesQuery,
297+
() -> valueFetcher,
298+
searchLookup,
299+
Lucene.STANDARD_ANALYZER,
300+
indexableFieldFunction,
301+
true
302+
);
303+
304+
// Index and Search
305+
try (Directory dir = newDirectory()) {
306+
IndexWriter iw = new IndexWriter(dir, new IndexWriterConfig(Lucene.STANDARD_ANALYZER));
307+
for (Document d : docs) {
308+
iw.addDocument(d);
309+
}
310+
try (IndexReader reader = DirectoryReader.open(iw)) {
311+
iw.close();
312+
IndexSearcher searcher = new IndexSearcher(reader);
313+
314+
// Rewrite should return the same query without attempting rewrite
315+
Query rewritten = derivedFieldQuery.rewrite(searcher);
316+
assertSame(derivedFieldQuery, rewritten);
317+
}
318+
}
319+
}
320+
321+
public void testNeedsRewriteWithRegularQuery() throws IOException {
322+
// Create lucene documents
323+
List<Document> docs = new ArrayList<>();
324+
for (String[] request : raw_requests) {
325+
Document document = new Document();
326+
document.add(new TextField("raw_request", request[0], Field.Store.YES));
327+
document.add(new KeywordField("status", request[1], Field.Store.YES));
328+
docs.add(document);
329+
}
330+
331+
// Mock SearchLookup
332+
SearchLookup searchLookup = mock(SearchLookup.class);
333+
SourceLookup sourceLookup = new SourceLookup();
334+
LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
335+
when(leafLookup.source()).thenReturn(sourceLookup);
336+
337+
// Mock DerivedFieldScript.Factory
338+
DerivedFieldScript.Factory factory = (params, lookup) -> (DerivedFieldScript.LeafFactory) ctx -> {
339+
when(searchLookup.getLeafSearchLookup(ctx)).thenReturn(leafLookup);
340+
return new DerivedFieldScript(params, lookup, ctx) {
341+
@Override
342+
public void execute() {
343+
addEmittedValue(raw_requests[sourceLookup.docId()][2]);
344+
}
345+
};
346+
};
347+
348+
// Create ValueFetcher from mocked DerivedFieldScript.Factory
349+
DerivedFieldScript.LeafFactory leafFactory = factory.newFactory((new Script("")).getParams(), searchLookup);
350+
Function<Object, IndexableField> indexableFieldFunction = DerivedFieldSupportedTypes.getIndexableFieldGeneratorType(
351+
"keyword",
352+
"ip_from_raw_request"
353+
);
354+
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(leafFactory, null);
355+
356+
// Create DerivedFieldQuery with regular TermQuery
357+
DerivedFieldQuery derivedFieldQuery = new DerivedFieldQuery(
358+
new TermQuery(new Term("ip_from_raw_request", "247.37.0.0")),
359+
() -> valueFetcher,
360+
searchLookup,
361+
Lucene.STANDARD_ANALYZER,
362+
indexableFieldFunction,
363+
true
364+
);
365+
366+
// Index and Search
367+
try (Directory dir = newDirectory()) {
368+
IndexWriter iw = new IndexWriter(dir, new IndexWriterConfig(Lucene.STANDARD_ANALYZER));
369+
for (Document d : docs) {
370+
iw.addDocument(d);
371+
}
372+
try (IndexReader reader = DirectoryReader.open(iw)) {
373+
iw.close();
374+
IndexSearcher searcher = new IndexSearcher(reader);
375+
376+
// Rewrite should perform rewrite on the inner query (TermQuery doesn't change in this case)
377+
Query rewritten = derivedFieldQuery.rewrite(searcher);
378+
// Since TermQuery doesn't rewrite to something different, we still get the same DerivedFieldQuery back
379+
assertSame(derivedFieldQuery, rewritten);
380+
}
381+
}
382+
}
184383
}

0 commit comments

Comments
 (0)