View Javadoc
1   package org.wikimedia.search.extra.regex;
2   
3   import java.io.IOException;
4   
5   import javax.annotation.Nullable;
6   
7   import org.apache.lucene.index.IndexReader;
8   import org.apache.lucene.index.LeafReaderContext;
9   import org.apache.lucene.search.ConstantScoreScorer;
10  import org.apache.lucene.search.ConstantScoreWeight;
11  import org.apache.lucene.search.IndexSearcher;
12  import org.apache.lucene.search.Query;
13  import org.apache.lucene.search.ScoreMode;
14  import org.apache.lucene.search.Scorer;
15  import org.apache.lucene.search.Weight;
16  import org.wikimedia.search.extra.regex.SourceRegexQuery.Rechecker;
17  import org.wikimedia.search.extra.regex.SourceRegexQueryBuilder.Settings;
18  import org.wikimedia.search.extra.util.FieldValues.Loader;
19  
20  import lombok.EqualsAndHashCode;
21  
22  /**
23   * Accelerated version of the source_regex query.
24   */
25  @EqualsAndHashCode(callSuper = true)
26  class AcceleratedSourceRegexQuery extends UnacceleratedSourceRegexQuery {
27      private final Query approximation;
28  
29      /**
30       * A new accelerated regex query.
31       *
32       * @param rechecker the rechecker used to perform the costly regex on doc content
33       * @param fieldPath the path to the field where the doc content is stored
34       * @param loader the loader used to load the field content
35       * @param settings the regex settings
36       * @param approximation the approximation query build over the trigram index
37       */
38      AcceleratedSourceRegexQuery(Rechecker rechecker, String fieldPath, Loader loader, Settings settings, Query approximation) {
39          super(rechecker, fieldPath, loader, settings);
40          this.approximation = approximation;
41      }
42  
43      @Override
44      public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
45          // Build the approximation based on trigrams
46          // Creating the Weight from the Searcher with needScore:false allows the searcher to cache our approximation.
47          final Weight approxWeight = searcher.createWeight(approximation, ScoreMode.COMPLETE_NO_SCORES, boost);
48          return new ConstantScoreWeight(this, 1F) {
49              @Override
50              public boolean isCacheable(LeafReaderContext leafReaderContext) {
51                  return false;
52              }
53  
54              @Override
55              @Nullable
56              public Scorer scorer(final LeafReaderContext context) throws IOException {
57                  final Scorer approxScorer = approxWeight.scorer(context);
58                  if (approxScorer == null) {
59                      return null;
60                  }
61                  return new ConstantScoreScorer(this, 1f, scoreMode, new RegexTwoPhaseIterator(approxScorer.iterator(), context));
62              }
63          };
64      }
65  
66      @Override
67      public Query rewrite(IndexReader reader) throws IOException {
68          Query approxRewritten = approximation.rewrite(reader);
69          if (approxRewritten != approximation) {
70              return new AcceleratedSourceRegexQuery(this.rechecker, this.fieldPath, this.loader, this.settings, approxRewritten);
71          }
72          return super.rewrite(reader);
73      }
74  
75      @Override
76      public String toString(String field) {
77          return "source_regex(accelerated):" + field;
78      }
79  }