1 package org.wikimedia.search.extra.termfreq;
2
3 import static org.wikimedia.search.extra.util.ConcreteIntPredicate.gt;
4 import static org.wikimedia.search.extra.util.ConcreteIntPredicate.gte;
5 import static org.wikimedia.search.extra.util.ConcreteIntPredicate.lt;
6 import static org.wikimedia.search.extra.util.ConcreteIntPredicate.lte;
7
8 import java.io.IOException;
9 import java.util.Objects;
10
11 import org.apache.lucene.index.Term;
12 import org.apache.lucene.search.Query;
13 import org.elasticsearch.common.ParseField;
14 import org.elasticsearch.common.ParsingException;
15 import org.elasticsearch.common.io.stream.StreamInput;
16 import org.elasticsearch.common.io.stream.StreamOutput;
17 import org.elasticsearch.common.xcontent.ObjectParser;
18 import org.elasticsearch.common.xcontent.XContentBuilder;
19 import org.elasticsearch.common.xcontent.XContentParser;
20 import org.elasticsearch.index.mapper.MappedFieldType;
21 import org.elasticsearch.index.query.AbstractQueryBuilder;
22 import org.elasticsearch.index.query.QueryShardContext;
23 import org.wikimedia.search.extra.util.ConcreteIntPredicate;
24
25 public class TermFreqFilterQueryBuilder extends AbstractQueryBuilder<TermFreqFilterQueryBuilder> {
26 public static final String NAME = "term_freq";
27
28 private static final ParseField FIELD = new ParseField("field");
29 private static final ParseField TERM = new ParseField("term");
30 private static final ParseField GT = new ParseField("gt");
31 private static final ParseField GTE = new ParseField("gte");
32 private static final ParseField LT = new ParseField("lt");
33 private static final ParseField LTE = new ParseField("lte");
34 private static final ParseField EQ = new ParseField("eq");
35
36 private static final ObjectParser<TermFreqFilterQueryBuilder, Void> PARSER = new ObjectParser<>(NAME, TermFreqFilterQueryBuilder::new);
37
38 static {
39 PARSER.declareString(TermFreqFilterQueryBuilder::setField, FIELD);
40 PARSER.declareString(TermFreqFilterQueryBuilder::setTerm, TERM);
41 PARSER.declareInt(TermFreqFilterQueryBuilder::setToStrict, LT);
42 PARSER.declareInt(TermFreqFilterQueryBuilder::setTo, LTE);
43 PARSER.declareInt(TermFreqFilterQueryBuilder::setFromStrict, GT);
44 PARSER.declareInt(TermFreqFilterQueryBuilder::setFrom, GTE);
45 PARSER.declareInt(TermFreqFilterQueryBuilder::setEqual, EQ);
46 AbstractQueryBuilder.declareStandardFields(PARSER);
47 }
48
49 private String field;
50 private String term;
51 private Integer to;
52 private boolean includeTo;
53 private Integer from;
54 private boolean includeFrom;
55 private Integer equal;
56
57 public TermFreqFilterQueryBuilder() {
58 }
59
60 public TermFreqFilterQueryBuilder(String field, String term) {
61 this.field = field;
62 this.term = term;
63 }
64
65 public TermFreqFilterQueryBuilder(StreamInput input) throws IOException {
66 super(input);
67 term = input.readString();
68 field = input.readString();
69 from = input.readOptionalVInt();
70 to = input.readOptionalVInt();
71 equal = input.readOptionalVInt();
72 includeFrom = input.readBoolean();
73 includeTo = input.readBoolean();
74 }
75
76 @Override
77 protected void doWriteTo(StreamOutput streamOutput) throws IOException {
78 streamOutput.writeString(term);
79 streamOutput.writeString(field);
80 streamOutput.writeOptionalVInt(this.from);
81 streamOutput.writeOptionalVInt(this.to);
82 streamOutput.writeOptionalVInt(this.equal);
83 streamOutput.writeBoolean(this.includeFrom);
84 streamOutput.writeBoolean(this.includeTo);
85 }
86
87 @SuppressWarnings("CyclomaticComplexity")
88 public static TermFreqFilterQueryBuilder fromXContent(XContentParser parser) throws IOException {
89 TermFreqFilterQueryBuilder builder;
90 try {
91 builder = PARSER.parse(parser, null);
92 } catch (IllegalArgumentException iae) {
93 throw new ParsingException(parser.getTokenLocation(), iae.getMessage(), iae);
94 }
95
96 if (builder.term == null) {
97 throw new ParsingException(parser.getTokenLocation(), TERM.getPreferredName() + " is mandatory");
98 }
99 if (builder.field == null) {
100 throw new ParsingException(parser.getTokenLocation(), FIELD.getPreferredName() + " is mandatory");
101 }
102 if (builder.equal != null) {
103 if (builder.from != null || builder.to != null) {
104 throw new ParsingException(parser.getTokenLocation(), EQ.getPreferredName() + " cannot be used with other comparators");
105 }
106 } else if (builder.from == null && builder.to == null) {
107 throw new ParsingException(parser.getTokenLocation(), "Invalid range provided eq or lt[e] and gt[e] must be provided");
108 }
109 if (builder.from != null && builder.to != null) {
110 checkRange(parser, builder);
111 }
112 return builder;
113 }
114
115 private static void checkRange(XContentParser parser, TermFreqFilterQueryBuilder builder) {
116 int minDiff = (builder.includeTo ? 0 : 1) + (builder.includeFrom ? 0 : 1);
117 int diff = builder.to - builder.from;
118 if (diff < minDiff) {
119 throw new ParsingException(parser.getTokenLocation(),
120 "Invalid range provided invalid range provided [" + builder.from + "," + builder.to + "]");
121 }
122 }
123
124 @Override
125 protected void doXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
126 xContentBuilder.startObject(NAME);
127 if (term != null) {
128 xContentBuilder.field(TERM.getPreferredName(), term);
129 }
130 if (field != null) {
131 xContentBuilder.field(FIELD.getPreferredName(), field);
132 }
133
134 if (equal != null) {
135 xContentBuilder.field(EQ.getPreferredName(), equal);
136 }
137 if (from != null) {
138 String gt = includeFrom ? GTE.getPreferredName() : GT.getPreferredName();
139 xContentBuilder.field(gt, from);
140 }
141 if (to != null) {
142 String lt = includeTo ? LTE.getPreferredName() : LT.getPreferredName();
143 xContentBuilder.field(lt, to);
144 }
145 printBoostAndQueryName(xContentBuilder);
146 xContentBuilder.endObject();
147 }
148
149 @Override
150 protected Query doToQuery(QueryShardContext queryShardContext) throws IOException {
151 MappedFieldType mapper = queryShardContext.fieldMapper(field);
152 if (mapper != null) {
153 if (!mapper.isSearchable()) {
154 throw new IllegalArgumentException("Cannot search on field [" + field + "] since it is not indexed.");
155 }
156 return new TermFreqFilterQuery(new Term(mapper.name(), term), buildPredicate());
157 }
158
159 return new TermFreqFilterQuery(new Term(field, term), buildPredicate());
160 }
161
162 private ConcreteIntPredicate buildPredicate() {
163 if (equal != null) {
164 return ConcreteIntPredicate.eq(equal);
165 }
166 ConcreteIntPredicate predicate = null;
167 if (from != null) {
168 predicate = buildFromPredicate();
169 }
170 if (to != null) {
171 predicate = predicate != null ? predicate.and(buildToPredicate()) : buildToPredicate();
172 }
173 if (predicate == null) {
174 throw new IllegalStateException("at least equal, from or to must be non null");
175 }
176 return predicate;
177 }
178
179 private ConcreteIntPredicate buildFromPredicate() {
180 if (includeFrom) {
181 return gte(from);
182 } else {
183 return gt(from);
184 }
185 }
186
187 private ConcreteIntPredicate buildToPredicate() {
188 if (includeTo) {
189 return lte(to);
190 } else {
191 return lt(to);
192 }
193 }
194
195 @Override
196 protected boolean doEquals(TermFreqFilterQueryBuilder termFreqQueryBuilder) {
197 return includeTo == termFreqQueryBuilder.includeTo &&
198 includeFrom == termFreqQueryBuilder.includeFrom &&
199 Objects.equals(termFreqQueryBuilder.from, from) &&
200 Objects.equals(termFreqQueryBuilder.to, to) &&
201 Objects.equals(termFreqQueryBuilder.equal, equal) &&
202 Objects.equals(termFreqQueryBuilder.term, term) &&
203 Objects.equals(termFreqQueryBuilder.field, field);
204 }
205
206 @Override
207 protected int doHashCode() {
208 return Objects.hash(field, term, includeFrom, from, includeTo, to, equal);
209 }
210
211 @Override
212 public String getWriteableName() {
213 return NAME;
214 }
215
216 public Integer getTo() {
217 return to;
218 }
219
220 public void setTo(Integer to) {
221 this.to = to;
222 this.includeTo = true;
223 }
224
225 public void setToStrict(Integer to) {
226 this.to = to;
227 this.includeTo = false;
228 }
229
230 public Integer getEqual() {
231 return equal;
232 }
233
234 public void setEqual(Integer equal) {
235 this.equal = equal;
236 }
237 public String getTerm() {
238 return term;
239 }
240
241 public void setTerm(String term) {
242 this.term = term;
243 }
244
245 public Integer getFrom() {
246 return from;
247 }
248
249 public void setFrom(Integer from) {
250 this.from = from;
251 this.includeFrom = true;
252 }
253
254 public void setFromStrict(Integer from) {
255 this.from = from;
256 this.includeFrom = false;
257 }
258
259 public boolean isIncludeTo() {
260 return includeTo;
261 }
262
263 public boolean isIncludeFrom() {
264 return includeFrom;
265 }
266
267 public String getField() {
268 return field;
269 }
270
271 public void setField(String field) {
272 this.field = field;
273 }
274 }