Replace deprecated Query#rewrite(IndexReader with rewrite(IndexSearcher

Query#rewrite(IndexReader reader) is deprecated in Lucene 9.7 and
removed in Lucene 10 in favor of Query#rewrite(IndexSearcher searcher).
The latter provides LeafCollector hooks to optimize for concurrent
queries. This commit cuts over usage of rewrite(IndexReader) to
rewrite(IndexSearcher) for Lucene 10 compatibility and upstream query
optimizations.

Signed-off-by: Nicholas Walter Knize <nknize@apache.org>

modules/parent-join/src/main/java/org/opensearch/join/query/HasChildQueryBuilder.java

@@ -416,11 +416,12 @@ public void visit(QueryVisitor visitor) {
         }
 
         @Override
-        public Query rewrite(IndexReader reader) throws IOException {
-            Query rewritten = super.rewrite(reader);
+        public Query rewrite(IndexSearcher searcher) throws IOException {
+            Query rewritten = super.rewrite(searcher);
             if (rewritten != this) {
                 return rewritten;
             }
+            IndexReader reader = searcher.getIndexReader();
             if (reader instanceof DirectoryReader) {
                 IndexSearcher indexSearcher = new IndexSearcher(reader);
                 indexSearcher.setQueryCache(null);

modules/percolator/src/main/java/org/opensearch/percolator/PercolateQuery.java

@@ -32,7 +32,6 @@
 
 package org.opensearch.percolator;
 
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.Explanation;
@@ -89,8 +88,8 @@ final class PercolateQuery extends Query implements Accountable {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query rewritten = candidateMatchesQuery.rewrite(reader);
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewritten = candidateMatchesQuery.rewrite(searcher);
         if (rewritten != candidateMatchesQuery) {
             return new PercolateQuery(
                 name,

modules/percolator/src/test/java/org/opensearch/percolator/CandidateQueryTests.java

@@ -1275,7 +1275,7 @@ private CustomQuery(Term term) {
         }
 
         @Override
-        public Query rewrite(IndexReader reader) throws IOException {
+        public Query rewrite(IndexSearcher searcher) throws IOException {
             return new TermQuery(term);
         }
 

server/src/main/java/org/apache/lucene/queries/BlendedTermQuery.java

@@ -42,6 +42,7 @@
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.DisjunctionMaxQuery;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.TermQuery;
@@ -93,11 +94,12 @@ public BlendedTermQuery(Term[] terms, float[] boosts) {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query rewritten = super.rewrite(reader);
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewritten = super.rewrite(searcher);
         if (rewritten != this) {
             return rewritten;
         }
+        IndexReader reader = searcher.getIndexReader();
         IndexReaderContext context = reader.getContext();
         TermStates[] ctx = new TermStates[terms.length];
         int[] docFreqs = new int[ctx.length];

server/src/main/java/org/opensearch/common/lucene/search/MoreLikeThisQuery.java

@@ -36,12 +36,12 @@
 import org.apache.lucene.analysis.TokenStream;
 import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
 import org.apache.lucene.index.Fields;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.index.Terms;
 import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.similarities.ClassicSimilarity;
@@ -144,12 +144,12 @@ public boolean equals(Object obj) {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query rewritten = super.rewrite(reader);
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewritten = super.rewrite(searcher);
         if (rewritten != this) {
             return rewritten;
         }
-        XMoreLikeThis mlt = new XMoreLikeThis(reader, similarity == null ? new ClassicSimilarity() : similarity);
+        XMoreLikeThis mlt = new XMoreLikeThis(searcher.getIndexReader(), similarity == null ? new ClassicSimilarity() : similarity);
 
         mlt.setFieldNames(moreLikeFields);
         mlt.setAnalyzer(analyzer);

server/src/main/java/org/opensearch/common/lucene/search/MultiPhrasePrefixQuery.java

@@ -39,6 +39,7 @@
 import org.apache.lucene.index.TermsEnum;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.MultiPhraseQuery;
 import org.apache.lucene.search.Query;
@@ -159,8 +160,8 @@ public int[] getPositions() {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query rewritten = super.rewrite(reader);
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewritten = super.rewrite(searcher);
         if (rewritten != this) {
             return rewritten;
         }
@@ -177,7 +178,7 @@ public Query rewrite(IndexReader reader) throws IOException {
         int position = positions.get(sizeMinus1);
         Set<Term> terms = new HashSet<>();
         for (Term term : suffixTerms) {
-            getPrefixTerms(terms, term, reader);
+            getPrefixTerms(terms, term, searcher.getIndexReader());
             if (terms.size() > maxExpansions) {
                 break;
             }

server/src/main/java/org/opensearch/common/lucene/search/function/FunctionScoreQuery.java

@@ -128,7 +128,7 @@ protected int doHashCode() {
 
         @Override
         protected ScoreFunction rewrite(IndexReader reader) throws IOException {
-            Query newFilter = filter.rewrite(reader);
+            Query newFilter = filter.rewrite(new IndexSearcher(reader));
             if (newFilter == filter) {
                 return this;
             }
@@ -322,16 +322,16 @@ public void visit(QueryVisitor visitor) {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query rewritten = super.rewrite(reader);
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query rewritten = super.rewrite(searcher);
         if (rewritten != this) {
             return rewritten;
         }
-        Query newQ = subQuery.rewrite(reader);
+        Query newQ = subQuery.rewrite(searcher);
         ScoreFunction[] newFunctions = new ScoreFunction[functions.length];
         boolean needsRewrite = (newQ != subQuery);
         for (int i = 0; i < functions.length; i++) {
-            newFunctions[i] = functions[i].rewrite(reader);
+            newFunctions[i] = functions[i].rewrite(searcher.getIndexReader());
             needsRewrite |= (newFunctions[i] != functions[i]);
         }
         if (needsRewrite) {

server/src/main/java/org/opensearch/common/lucene/search/function/ScriptScoreQuery.java

@@ -32,7 +32,6 @@
 
 package org.opensearch.common.lucene.search.function;
 
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.DocIdSetIterator;
@@ -105,12 +104,12 @@ public ScriptScoreQuery(
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query newQ = subQuery.rewrite(reader);
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query newQ = subQuery.rewrite(searcher);
         if (newQ != subQuery) {
             return new ScriptScoreQuery(newQ, queryName, script, scriptBuilder, minScore, indexName, shardId, indexVersion);
         }
-        return super.rewrite(reader);
+        return super.rewrite(searcher);
     }
 
     @Override

server/src/main/java/org/opensearch/index/document/SortedUnsignedLongDocValuesRangeQuery.java

@@ -10,7 +10,6 @@
 
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.FieldInfo;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
@@ -87,12 +86,12 @@ public String toString(String field) {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
+    public Query rewrite(IndexSearcher searcher) throws IOException {
         if (Long.compareUnsigned(lowerValue, Numbers.MIN_UNSIGNED_LONG_VALUE_AS_LONG) == 0
             && Long.compareUnsigned(upperValue, Numbers.MAX_UNSIGNED_LONG_VALUE_AS_LONG) == 0) {
             return new FieldExistsQuery(field);
         }
-        return super.rewrite(reader);
+        return super.rewrite(searcher);
     }
 
     abstract SortedNumericDocValues getValues(LeafReader reader, String field) throws IOException;

server/src/main/java/org/opensearch/index/query/DateRangeIncludingNowQuery.java

@@ -32,8 +32,8 @@
 
 package org.opensearch.index.query;
 
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 
@@ -60,7 +60,7 @@ public Query getQuery() {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
+    public Query rewrite(IndexSearcher searcher) throws IOException {
         return in;
     }
 

server/src/main/java/org/opensearch/index/search/OpenSearchToParentBlockJoinQuery.java

@@ -32,7 +32,6 @@
 
 package org.opensearch.index.search;
 
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
@@ -82,8 +81,8 @@ public ScoreMode getScoreMode() {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query innerRewrite = query.rewrite(reader);
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        Query innerRewrite = query.rewrite(searcher);
         if (innerRewrite != query) {
             // Right now ToParentBlockJoinQuery always rewrites to a ToParentBlockJoinQuery
             // so the else block will never be used. It is useful in the case that
@@ -97,7 +96,7 @@ public Query rewrite(IndexReader reader) throws IOException {
                 return innerRewrite;
             }
         }
-        return super.rewrite(reader);
+        return super.rewrite(searcher);
     }
 
     @Override

server/src/main/java/org/opensearch/lucene/queries/MinDocQuery.java

@@ -86,7 +86,8 @@ public boolean equals(Object obj) {
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
+    public Query rewrite(IndexSearcher searcher) throws IOException {
+        final IndexReader reader = searcher.getIndexReader();
         if (Objects.equals(reader.getContext().id(), readerId) == false) {
             return new MinDocQuery(minDoc, reader.getContext().id());
         }

server/src/main/java/org/opensearch/search/suggest/completion/CompletionSuggester.java

@@ -104,7 +104,7 @@ protected Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Sugges
     }
 
     private static void suggest(IndexSearcher searcher, CompletionQuery query, TopSuggestDocsCollector collector) throws IOException {
-        query = (CompletionQuery) query.rewrite(searcher.getIndexReader());
+        query = (CompletionQuery) query.rewrite(searcher);
         Weight weight = query.createWeight(searcher, collector.scoreMode(), 1f);
         for (LeafReaderContext context : searcher.getIndexReader().leaves()) {
             BulkScorer scorer = weight.bulkScorer(context);
@@ -124,6 +124,7 @@ private static void suggest(IndexSearcher searcher, CompletionQuery query, TopSu
                 }
             }
         }
+        collector.finish();
     }
 
     @Override

server/src/test/java/org/opensearch/deps/lucene/VectorHighlighterTests.java

@@ -125,7 +125,7 @@ public void testVectorHighlighterPrefixQuery() throws Exception {
         assertThat(fragment, nullValue());
 
         prefixQuery = new PrefixQuery(new Term("content", "ba"), PrefixQuery.SCORING_BOOLEAN_REWRITE);
-        Query rewriteQuery = prefixQuery.rewrite(reader);
+        Query rewriteQuery = prefixQuery.rewrite(searcher);
         fragment = highlighter.getBestFragment(highlighter.getFieldQuery(rewriteQuery), reader, topDocs.scoreDocs[0].doc, "content", 30);
         assertThat(fragment, notNullValue());
 

server/src/test/java/org/opensearch/index/mapper/DateFieldTypeTests.java

@@ -40,6 +40,7 @@
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.IndexSortSortedNumericDocValuesRangeQuery;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.IndexOrDocValuesQuery;
@@ -259,7 +260,10 @@ public void testRangeQuery() throws IOException {
             LongPoint.newRangeQuery("field", instant1, instant2),
             SortedNumericDocValuesField.newSlowRangeQuery("field", instant1, instant2)
         );
-        assertEquals(expected, ft.rangeQuery(date1, date2, true, true, null, null, null, context).rewrite(new MultiReader()));
+        assertEquals(
+            expected,
+            ft.rangeQuery(date1, date2, true, true, null, null, null, context).rewrite(new IndexSearcher(new MultiReader()))
+        );
 
         instant1 = nowInMillis;
         instant2 = instant1 + 100;

server/src/test/java/org/opensearch/index/query/SpanMultiTermQueryBuilderTests.java

@@ -284,8 +284,9 @@ public void testTermExpansionExceptionOnSpanFailure() throws Exception {
                     BooleanQuery.setMaxClauseCount(1);
                     try {
                         QueryBuilder queryBuilder = new SpanMultiTermQueryBuilder(QueryBuilders.prefixQuery("body", "bar"));
-                        Query query = queryBuilder.toQuery(createShardContext(new IndexSearcher(reader)));
-                        RuntimeException exc = expectThrows(RuntimeException.class, () -> query.rewrite(reader));
+                        IndexSearcher searcher = new IndexSearcher(reader);
+                        Query query = queryBuilder.toQuery(createShardContext(searcher));
+                        RuntimeException exc = expectThrows(RuntimeException.class, () -> query.rewrite(searcher));
                         assertThat(exc.getMessage(), containsString("maxClauseCount"));
                     } finally {
                         BooleanQuery.setMaxClauseCount(origBoolMaxClauseCount);

server/src/test/java/org/opensearch/index/search/OpenSearchToParentBlockJoinQueryTests.java

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.index.MultiReader;
 import org.apache.lucene.index.Term;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.PhraseQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
@@ -112,7 +113,7 @@ public void testRewrite() throws IOException {
             ScoreMode.Avg,
             "nested"
         );
-        Query rewritten = q.rewrite(new MultiReader());
+        Query rewritten = q.rewrite(new IndexSearcher(new MultiReader()));
         assertEquals(expected, rewritten);
     }
 }

server/src/test/java/org/opensearch/lucene/queries/MinDocQueryTests.java

@@ -61,10 +61,11 @@ public void testBasics() {
 
     public void testRewrite() throws Exception {
         IndexReader reader = new MultiReader();
+        IndexSearcher searcher = new IndexSearcher(reader);
         MinDocQuery query = new MinDocQuery(42);
-        Query rewritten = query.rewrite(reader);
+        Query rewritten = query.rewrite(searcher);
         QueryUtils.checkUnequal(query, rewritten);
-        Query rewritten2 = rewritten.rewrite(reader);
+        Query rewritten2 = rewritten.rewrite(searcher);
         assertSame(rewritten, rewritten2);
     }
 

server/src/test/java/org/opensearch/lucene/queries/SpanMatchNoDocsQueryTests.java

@@ -56,19 +56,11 @@
 
 public class SpanMatchNoDocsQueryTests extends OpenSearchTestCase {
     public void testSimple() throws Exception {
-        Directory dir = newDirectory();
-        IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig());
-        IndexReader ir = DirectoryReader.open(iw);
-
         SpanMatchNoDocsQuery query = new SpanMatchNoDocsQuery("field", "a good reason");
         assertEquals(query.toString(), "SpanMatchNoDocsQuery(\"a good reason\")");
-        Query rewrite = query.rewrite(ir);
+        Query rewrite = query.rewrite((IndexSearcher) null);
         assertTrue(rewrite instanceof SpanMatchNoDocsQuery);
         assertEquals(rewrite.toString(), "SpanMatchNoDocsQuery(\"a good reason\")");
-
-        iw.close();
-        ir.close();
-        dir.close();
     }
 
     public void testQuery() throws Exception {
@@ -101,7 +93,7 @@ public void testQuery() throws Exception {
         assertEquals(searcher.count(orQuery), 1);
         hits = searcher.search(orQuery, 1000).scoreDocs;
         assertEquals(1, hits.length);
-        Query rewrite = orQuery.rewrite(ir);
+        Query rewrite = orQuery.rewrite(searcher);
         assertEquals(rewrite, orQuery);
 
         SpanNearQuery nearQuery = new SpanNearQuery(
@@ -112,7 +104,7 @@ public void testQuery() throws Exception {
         assertEquals(searcher.count(nearQuery), 0);
         hits = searcher.search(nearQuery, 1000).scoreDocs;
         assertEquals(0, hits.length);
-        rewrite = nearQuery.rewrite(ir);
+        rewrite = nearQuery.rewrite(searcher);
         assertEquals(rewrite, nearQuery);
 
         iw.close();