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

Replace Query#rewrite(IndexReader) with Query#rewrite(IndexSearcher)
benwtrent · Mar 9, 2023 · 17cac27 · 17cac27
1 parent 5ff8ea4
commit 17cac27
Show file tree

Hide file tree

Showing 106 changed files with 338 additions and 343 deletions.
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
@@ -18,6 +18,9 @@ API Changes
 * GITHUB#12173: TermInSetQuery#getTermData has been deprecated. This exposes internal implementation details that we
   may want to change in the future, and users shouldn't rely on the encoding directly. (Greg Miller)
 
+* 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
@@ -21,7 +21,6 @@
 import java.util.Arrays;
 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.search.ConstantScoreScorer;
@@ -90,8 +89,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.Term;
@@ -52,12 +51,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.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
@@ -84,11 +83,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);
   }
 
   @Override

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.LeafReaderContext;
 import org.apache.lucene.index.SortedDocValues;
 import org.apache.lucene.index.SortedSetDocValues;
@@ -98,11 +97,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);
   }
 
   @Override

diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java
@@ -60,12 +60,12 @@ public AbstractKnnVectorQuery(String field, int k, Query filter) {
   }
 
   @Override
-  public Query rewrite(IndexReader reader) throws IOException {
+  public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+    IndexReader reader = indexSearcher.getIndexReader();
     TopDocs[] perLeafResults = new TopDocs[reader.leaves().size()];
 
     Weight filterWeight = null;
     if (filter != null) {
-      IndexSearcher indexSearcher = new IndexSearcher(reader);
       BooleanQuery booleanQuery =
           new BooleanQuery.Builder()
               .add(filter, BooleanClause.Occur.FILTER)

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;
 
 /**
@@ -247,7 +246,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");
     }
@@ -286,12 +285,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
@@ -566,7 +565,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.
@@ -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