GITHUB-11838 Change API to allow concurrent query rewrite (#11840)

Replace Query#rewrite(IndexReader) with Query#rewrite(IndexSearcher)
apache · Oct 19, 2022 · 6cde41c · 6cde41c
1 parent 05971b3
commit 6cde41c
Show file tree

Hide file tree

Showing 107 changed files with 344 additions and 350 deletions.
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
@@ -34,6 +34,9 @@ API Changes
 * GITHUB#11813: Remove Operations.isFinite: the recursive implementation could be problematic
   for large automatons (WildcardQuery, PrefixQuery, RegExpQuery, etc).  (taroplus, Robert Muir)
 
+* GITHUB#11840: Query rewrite now takes an IndexSearcher instead of IndexReader to enable concurrent
+  rewriting. (Patrick Zhai)
+
 New Features
 ---------------------
 

diff --git a/lucene/classification/src/java/org/apache/lucene/classification/utils/NearestFuzzyQuery.java b/lucene/classification/src/java/org/apache/lucene/classification/utils/NearestFuzzyQuery.java
@@ -35,6 +35,7 @@
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.FuzzyTermsEnum;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 import org.apache.lucene.search.TermQuery;
@@ -214,7 +215,8 @@ private Query newTermQuery(IndexReader reader, Term term) throws IOException {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    IndexReader reader = indexSearcher.getIndexReader();
     ScoreTermQueue q = new ScoreTermQueue(MAX_NUM_TERMS);
     // load up the list of possible terms
     for (FieldVals f : fieldVals) {

diff --git a/lucene/core/src/java/org/apache/lucene/document/BinaryRangeFieldRangeQuery.java b/lucene/core/src/java/org/apache/lucene/document/BinaryRangeFieldRangeQuery.java
@@ -22,7 +22,6 @@
 import java.util.Objects;
 import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.DocValues;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.ConstantScoreScorer;
@@ -87,8 +86,8 @@ public void visit(QueryVisitor visitor) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    return super.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    return super.rewrite(indexSearcher);
   }
 
   private BinaryRangeDocValues getValues(LeafReader reader, String field) throws IOException {

diff --git a/lucene/core/src/java/org/apache/lucene/document/DoubleRangeSlowRangeQuery.java b/lucene/core/src/java/org/apache/lucene/document/DoubleRangeSlowRangeQuery.java
@@ -20,7 +20,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Objects;
-import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 
@@ -79,8 +79,8 @@ public String toString(String field) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    return super.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    return super.rewrite(indexSearcher);
   }
 
   private static byte[] encodeRanges(double[] min, double[] max) {

diff --git a/lucene/core/src/java/org/apache/lucene/document/FeatureField.java b/lucene/core/src/java/org/apache/lucene/document/FeatureField.java
@@ -31,6 +31,7 @@
 import org.apache.lucene.search.DoubleValuesSource;
 import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.FieldDoc;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.similarities.BM25Similarity;
@@ -223,7 +224,7 @@ abstract static class FeatureFunction {
 
     abstract Explanation explain(String field, String feature, float w, int freq);
 
-    FeatureFunction rewrite(IndexReader reader) throws IOException {
+    FeatureFunction rewrite(IndexSearcher indexSearcher) throws IOException {
       return this;
     }
   }
@@ -340,11 +341,11 @@ static final class SaturationFunction extends FeatureFunction {
     }
 
     @Override
-    public FeatureFunction rewrite(IndexReader reader) throws IOException {
+    public FeatureFunction rewrite(IndexSearcher indexSearcher) throws IOException {
       if (pivot != null) {
-        return super.rewrite(reader);
+        return super.rewrite(indexSearcher);
       }
-      float newPivot = computePivotFeatureValue(reader, field, feature);
+      float newPivot = computePivotFeatureValue(indexSearcher.getIndexReader(), field, feature);
       return new SaturationFunction(field, feature, newPivot);
     }
 

diff --git a/lucene/core/src/java/org/apache/lucene/document/FeatureQuery.java b/lucene/core/src/java/org/apache/lucene/document/FeatureQuery.java
@@ -20,7 +20,6 @@
 import java.util.Objects;
 import org.apache.lucene.document.FeatureField.FeatureFunction;
 import org.apache.lucene.index.ImpactsEnum;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.PostingsEnum;
 import org.apache.lucene.index.Terms;
@@ -50,12 +49,12 @@ final class FeatureQuery extends Query {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    FeatureFunction rewritten = function.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    FeatureFunction rewritten = function.rewrite(indexSearcher);
     if (function != rewritten) {
       return new FeatureQuery(fieldName, featureName, rewritten);
     }
-    return super.rewrite(reader);
+    return super.rewrite(indexSearcher);
   }
 
   @Override

diff --git a/lucene/core/src/java/org/apache/lucene/document/FloatRangeSlowRangeQuery.java b/lucene/core/src/java/org/apache/lucene/document/FloatRangeSlowRangeQuery.java
@@ -20,7 +20,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Objects;
-import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 
@@ -79,8 +79,8 @@ public String toString(String field) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    return super.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    return super.rewrite(indexSearcher);
   }
 
   private static byte[] encodeRanges(float[] min, float[] max) {

diff --git a/lucene/core/src/java/org/apache/lucene/document/IntRangeSlowRangeQuery.java b/lucene/core/src/java/org/apache/lucene/document/IntRangeSlowRangeQuery.java
@@ -19,7 +19,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Objects;
-import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 
@@ -77,8 +77,8 @@ public String toString(String field) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    return super.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    return super.rewrite(indexSearcher);
   }
 
   private static byte[] encodeRanges(int[] min, int[] max) {

diff --git a/lucene/core/src/java/org/apache/lucene/document/LongRangeSlowRangeQuery.java b/lucene/core/src/java/org/apache/lucene/document/LongRangeSlowRangeQuery.java
@@ -20,7 +20,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Objects;
-import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
 
@@ -79,8 +79,8 @@ public String toString(String field) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    return super.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    return super.rewrite(indexSearcher);
   }
 
   private static byte[] encodeRanges(long[] min, long[] max) {

diff --git a/lucene/core/src/java/org/apache/lucene/document/SortedNumericDocValuesRangeQuery.java b/lucene/core/src/java/org/apache/lucene/document/SortedNumericDocValuesRangeQuery.java
@@ -19,7 +19,6 @@
 import java.io.IOException;
 import java.util.Objects;
 import org.apache.lucene.index.DocValues;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
@@ -89,11 +88,11 @@ public String toString(String field) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
     if (lowerValue == Long.MIN_VALUE && upperValue == Long.MAX_VALUE) {
       return new FieldExistsQuery(field);
     }
-    return super.rewrite(reader);
+    return super.rewrite(indexSearcher);
   }
 
   abstract SortedNumericDocValues getValues(LeafReader reader, String field) throws IOException;

diff --git a/lucene/core/src/java/org/apache/lucene/document/SortedSetDocValuesRangeQuery.java b/lucene/core/src/java/org/apache/lucene/document/SortedSetDocValuesRangeQuery.java
@@ -19,7 +19,6 @@
 import java.io.IOException;
 import java.util.Objects;
 import org.apache.lucene.index.DocValues;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedDocValues;
@@ -103,11 +102,11 @@ public String toString(String field) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
     if (lowerValue == null && upperValue == null) {
       return new FieldExistsQuery(field);
     }
-    return super.rewrite(reader);
+    return super.rewrite(indexSearcher);
   }
 
   abstract SortedSetDocValues getValues(LeafReader reader, String field) throws IOException;

diff --git a/lucene/core/src/java/org/apache/lucene/search/BlendedTermQuery.java b/lucene/core/src/java/org/apache/lucene/search/BlendedTermQuery.java
@@ -19,7 +19,6 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
@@ -268,11 +267,12 @@ public String toString(String field) {
   }
 
   @Override
-  public final Query rewrite(IndexReader reader) throws IOException {
+  public final Query rewrite(IndexSearcher indexSearcher) throws IOException {
     final TermStates[] contexts = ArrayUtil.copyOfSubArray(this.contexts, 0, this.contexts.length);
     for (int i = 0; i < contexts.length; ++i) {
-      if (contexts[i] == null || contexts[i].wasBuiltFor(reader.getContext()) == false) {
-        contexts[i] = TermStates.build(reader.getContext(), terms[i], true);
+      if (contexts[i] == null
+          || contexts[i].wasBuiltFor(indexSearcher.getTopReaderContext()) == false) {
+        contexts[i] = TermStates.build(indexSearcher.getTopReaderContext(), terms[i], true);
       }
     }
 
@@ -287,7 +287,7 @@ public final Query rewrite(IndexReader reader) throws IOException {
     }
 
     for (int i = 0; i < contexts.length; ++i) {
-      contexts[i] = adjustFrequencies(reader.getContext(), contexts[i], df, ttf);
+      contexts[i] = adjustFrequencies(indexSearcher.getTopReaderContext(), contexts[i], df, ttf);
     }
 
     Query[] termQueries = new Query[terms.length];

diff --git a/lucene/core/src/java/org/apache/lucene/search/BooleanQuery.java b/lucene/core/src/java/org/apache/lucene/search/BooleanQuery.java
@@ -30,7 +30,6 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.Predicate;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.BooleanClause.Occur;
 
 /**
@@ -193,7 +192,7 @@ public final Iterator<BooleanClause> iterator() {
 
   // Utility method for rewriting BooleanQuery when scores are not needed.
   // This is called from ConstantScoreQuery#rewrite
-  BooleanQuery rewriteNoScoring(IndexReader reader) throws IOException {
+  BooleanQuery rewriteNoScoring(IndexSearcher indexSearcher) throws IOException {
     boolean actuallyRewritten = false;
     BooleanQuery.Builder newQuery =
         new BooleanQuery.Builder().setMinimumNumberShouldMatch(getMinimumNumberShouldMatch());
@@ -204,7 +203,7 @@ BooleanQuery rewriteNoScoring(IndexReader reader) throws IOException {
 
     for (BooleanClause clause : clauses) {
       Query query = clause.getQuery();
-      Query rewritten = new ConstantScoreQuery(query).rewrite(reader);
+      Query rewritten = new ConstantScoreQuery(query).rewrite(indexSearcher);
       if (rewritten instanceof ConstantScoreQuery) {
         rewritten = ((ConstantScoreQuery) rewritten).getQuery();
       }
@@ -238,7 +237,7 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
     if (clauses.size() == 0) {
       return new MatchNoDocsQuery("empty BooleanQuery");
     }
@@ -277,12 +276,12 @@ public Query rewrite(IndexReader reader) throws IOException {
         Query rewritten;
         if (occur == Occur.FILTER || occur == Occur.MUST_NOT) {
           // Clauses that are not involved in scoring can get some extra simplifications
-          rewritten = new ConstantScoreQuery(query).rewrite(reader);
+          rewritten = new ConstantScoreQuery(query).rewrite(indexSearcher);
           if (rewritten instanceof ConstantScoreQuery) {
             rewritten = ((ConstantScoreQuery) rewritten).getQuery();
           }
         } else {
-          rewritten = query.rewrite(reader);
+          rewritten = query.rewrite(indexSearcher);
         }
         if (rewritten != query || query.getClass() == MatchNoDocsQuery.class) {
           // rewrite clause
@@ -557,7 +556,7 @@ public Query rewrite(IndexReader reader) throws IOException {
       }
     }
 
-    return super.rewrite(reader);
+    return super.rewrite(indexSearcher);
   }
 
   @Override

diff --git a/lucene/core/src/java/org/apache/lucene/search/BoostQuery.java b/lucene/core/src/java/org/apache/lucene/search/BoostQuery.java
@@ -18,7 +18,6 @@
 
 import java.io.IOException;
 import java.util.Objects;
-import org.apache.lucene.index.IndexReader;
 
 /**
  * A {@link Query} wrapper that allows to give a boost to the wrapped query. Boost values that are
@@ -73,8 +72,8 @@ public int hashCode() {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    final Query rewritten = query.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    final Query rewritten = query.rewrite(indexSearcher);
 
     if (boost == 1f) {
       return rewritten;
@@ -99,7 +98,7 @@ public Query rewrite(IndexReader reader) throws IOException {
       return new BoostQuery(rewritten, boost);
     }
 
-    return super.rewrite(reader);
+    return super.rewrite(indexSearcher);
   }
 
   @Override

diff --git a/lucene/core/src/java/org/apache/lucene/search/ConstantScoreQuery.java b/lucene/core/src/java/org/apache/lucene/search/ConstantScoreQuery.java
@@ -18,7 +18,6 @@
 
 import java.io.IOException;
 import java.util.Objects;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.util.Bits;
 
@@ -40,8 +39,9 @@ public Query getQuery() {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
-    Query rewritten = query.rewrite(reader);
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+
+    Query rewritten = query.rewrite(indexSearcher);
 
     // Do some extra simplifications that are legal since scores are not needed on the wrapped
     // query.
@@ -50,7 +50,7 @@ public Query rewrite(IndexReader reader) throws IOException {
     } else if (rewritten instanceof ConstantScoreQuery) {
       rewritten = ((ConstantScoreQuery) rewritten).getQuery();
     } else if (rewritten instanceof BooleanQuery) {
-      rewritten = ((BooleanQuery) rewritten).rewriteNoScoring(reader);
+      rewritten = ((BooleanQuery) rewritten).rewriteNoScoring(indexSearcher);
     }
 
     if (rewritten.getClass() == MatchNoDocsQuery.class) {
@@ -70,7 +70,7 @@ public Query rewrite(IndexReader reader) throws IOException {
       return new ConstantScoreQuery(((BoostQuery) rewritten).getQuery());
     }
 
-    return super.rewrite(reader);
+    return super.rewrite(indexSearcher);
   }
 
   @Override