1919
2020package org .elasticsearch .index .query ;
2121
22+ import com .google .common .collect .Iterables ;
23+
24+ import org .apache .lucene .queries .TermsQuery ;
25+ import org .apache .lucene .search .Query ;
26+ import org .apache .lucene .util .BytesRef ;
27+ import org .elasticsearch .common .io .stream .StreamInput ;
28+ import org .elasticsearch .common .io .stream .StreamOutput ;
29+ import org .elasticsearch .common .io .stream .Streamable ;
30+ import org .elasticsearch .common .lucene .search .Queries ;
2231import org .elasticsearch .common .xcontent .XContentBuilder ;
32+ import org .elasticsearch .index .mapper .Uid ;
33+ import org .elasticsearch .index .mapper .internal .UidFieldMapper ;
2334
2435import java .io .IOException ;
2536import java .util .ArrayList ;
2637import java .util .Arrays ;
27- import java .util .List ;
38+ import java .util .Collection ;
39+ import java .util .Objects ;
2840
2941/**
3042 * A query that will return only documents matching specific ids (and a type).
3143 */
32- public class IdsQueryBuilder extends BaseQueryBuilder implements BoostableQueryBuilder <IdsQueryBuilder > {
44+ public class IdsQueryBuilder extends BaseQueryBuilder implements Streamable , BoostableQueryBuilder <IdsQueryBuilder > {
3345
34- private final List <String > types ;
46+ private Collection <String > types = new ArrayList <>() ;
3547
36- private List <String > values = new ArrayList <>();
48+ private Collection <String > ids = new ArrayList <>();
3749
38- private float boost = - 1 ;
50+ private float boost = 1.0f ;
3951
4052 private String queryName ;
4153
4254 public IdsQueryBuilder (String ... types ) {
43- this .types = types == null ? null : Arrays .asList (types );
55+ this .types = ( types == null || types . length == 0 ) ? new ArrayList < String >() : Arrays .asList (types );
4456 }
4557
4658 /**
47- * Adds ids to the filter.
59+ * Get the types used in this query
60+ * @return the types
61+ */
62+ public Collection <String > types () {
63+ return this .types ;
64+ }
65+
66+ /**
67+ * Adds ids to the query.
4868 */
4969 public IdsQueryBuilder addIds (String ... ids ) {
50- values .addAll (Arrays .asList (ids ));
70+ this . ids .addAll (Arrays .asList (ids ));
5171 return this ;
5272 }
5373
5474 /**
55- * Adds ids to the filter .
75+ * Adds ids to the query .
5676 */
5777 public IdsQueryBuilder ids (String ... ids ) {
5878 return addIds (ids );
5979 }
6080
81+ /**
82+ * Gets the ids for the query.
83+ */
84+ public Collection <String > ids () {
85+ return this .ids ;
86+ }
87+
6188 /**
6289 * Sets the boost for this query. Documents matching this query will (in addition to the normal
6390 * weightings) have their score multiplied by the boost provided.
@@ -69,19 +96,33 @@ public IdsQueryBuilder boost(float boost) {
6996 }
7097
7198 /**
72- * Sets the query name for the filter that can be used when searching for matched_filters per hit.
99+ * Gets the boost for this query.
100+ */
101+ public float boost () {
102+ return this .boost ;
103+ }
104+
105+ /**
106+ * Sets the query name for the query that can be used when searching for matched_filters per hit.
73107 */
74108 public IdsQueryBuilder queryName (String queryName ) {
75109 this .queryName = queryName ;
76110 return this ;
77111 }
78112
113+ /**
114+ * Gets the query name for the query.
115+ */
116+ public String queryName () {
117+ return this .queryName ;
118+ }
119+
79120 @ Override
80121 protected void doXContent (XContentBuilder builder , Params params ) throws IOException {
81122 builder .startObject (IdsQueryParser .NAME );
82123 if (types != null ) {
83124 if (types .size () == 1 ) {
84- builder .field ("type" , types .get ( 0 ));
125+ builder .field ("type" , types .iterator (). next ( ));
85126 } else {
86127 builder .startArray ("types" );
87128 for (Object type : types ) {
@@ -91,7 +132,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
91132 }
92133 }
93134 builder .startArray ("values" );
94- for (Object value : values ) {
135+ for (Object value : ids ) {
95136 builder .value (value );
96137 }
97138 builder .endArray ();
@@ -108,4 +149,84 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
108149 protected String parserName () {
109150 return IdsQueryParser .NAME ;
110151 }
152+
153+ public Query toQuery (QueryParseContext parseContext ) throws IOException , QueryParsingException {
154+ ArrayList <BytesRef > ids = new ArrayList <>();
155+
156+ for (String value : this .ids ) {
157+ BytesRef ref = new BytesRef (value );
158+ ids .add (ref );
159+ }
160+
161+ if (ids .isEmpty ()) {
162+ return Queries .newMatchNoDocsQuery ();
163+ }
164+
165+ if (types == null || types .isEmpty ()) {
166+ types = parseContext .queryTypes ();
167+ } else if (types .size () == 1 && Iterables .getFirst (types , null ).equals ("_all" )) {
168+ types = parseContext .mapperService ().types ();
169+ }
170+
171+ TermsQuery query = new TermsQuery (UidFieldMapper .NAME , Uid .createTypeUids (types , ids ));
172+ query .setBoost (boost );
173+ if (queryName != null ) {
174+ parseContext .addNamedQuery (queryName , query );
175+ }
176+ return query ;
177+ }
178+
179+ @ Override
180+ public QueryValidationException validate () {
181+ // all fields can be empty or null
182+ return null ;
183+ }
184+
185+ @ Override
186+ public void readFrom (StreamInput in ) throws IOException {
187+ int typeSize = in .readInt ();
188+ for (int i = 0 ; i < typeSize ; i ++) {
189+ types .add (in .readString ());
190+ }
191+ int valueSize = in .readInt ();
192+ for (int i = 0 ; i < valueSize ; i ++) {
193+ ids .add (in .readString ());
194+ }
195+ queryName = in .readOptionalString ();
196+ boost = in .readFloat ();
197+ }
198+
199+ @ Override
200+ public void writeTo (StreamOutput out ) throws IOException {
201+ out .writeInt (types .size ());
202+ for (String type : types ) {
203+ out .writeString (type );
204+ }
205+ out .writeInt (ids .size ());
206+ for (String value : ids ) {
207+ out .writeString (value );
208+ }
209+ out .writeOptionalString (queryName );
210+ out .writeFloat (boost );
211+ }
212+
213+ @ Override
214+ public int hashCode () {
215+ return Objects .hash (ids , types , boost , queryName );
216+ }
217+
218+ @ Override
219+ public boolean equals (Object obj ) {
220+ if (this == obj ) {
221+ return true ;
222+ }
223+ if (obj == null || getClass () != obj .getClass ()) {
224+ return false ;
225+ }
226+ IdsQueryBuilder other = (IdsQueryBuilder ) obj ;
227+ return Objects .equals (ids , other .ids ) &&
228+ Objects .equals (types , other .types ) &&
229+ Objects .equals (boost , other .boost ) &&
230+ Objects .equals (queryName , other .queryName );
231+ }
111232}
0 commit comments