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
24
25 @EqualsAndHashCode(callSuper = true)
26 class AcceleratedSourceRegexQuery extends UnacceleratedSourceRegexQuery {
27 private final Query approximation;
28
29
30
31
32
33
34
35
36
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
46
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 }