Skip to content

Commit 3cefd40

Browse files
jimczijimczi
committed
Allow terms query in _rollup_search (#30973)
This change adds the `terms` query to the list of accepted queries for the _rollup_search endpoint.
1 parent f485e8c commit 3cefd40

File tree

2 files changed

+110
-83
lines changed

2 files changed

+110
-83
lines changed

‎x-pack/plugin/rollup/src/main/java/org/elasticsearch/xpack/rollup/action/TransportRollupSearchAction.java

Lines changed: 89 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import org.elasticsearch.index.query.QueryBuilder;
4040
import org.elasticsearch.index.query.RangeQueryBuilder;
4141
import org.elasticsearch.index.query.TermQueryBuilder;
42+
import org.elasticsearch.index.query.TermsQueryBuilder;
4243
import org.elasticsearch.script.ScriptService;
4344
import org.elasticsearch.search.aggregations.AggregationBuilder;
4445
import org.elasticsearch.search.aggregations.AggregatorFactories;
@@ -66,6 +67,7 @@
6667
import java.io.IOException;
6768
import java.util.ArrayList;
6869
import java.util.Arrays;
70+
import java.util.Collections;
6971
import java.util.HashSet;
7072
import java.util.List;
7173
import java.util.Objects;
@@ -276,91 +278,38 @@ static QueryBuilder rewriteQuery(QueryBuilder builder, Set<RollupJobCaps> jobCap
276278
rewriteQuery(((BoostingQueryBuilder)builder).positiveQuery(), jobCaps));
277279
} else if (builder.getWriteableName().equals(DisMaxQueryBuilder.NAME)) {
278280
DisMaxQueryBuilder rewritten = new DisMaxQueryBuilder();
279-
((DisMaxQueryBuilder)builder).innerQueries().forEach(query -> rewritten.add(rewriteQuery(query, jobCaps)));
281+
((DisMaxQueryBuilder) builder).innerQueries().forEach(query -> rewritten.add(rewriteQuery(query, jobCaps)));
280282
return rewritten;
281-
} else if (builder.getWriteableName().equals(RangeQueryBuilder.NAME) || builder.getWriteableName().equals(TermQueryBuilder.NAME)) {
282-
283-
String fieldName = builder.getWriteableName().equals(RangeQueryBuilder.NAME)
284-
? ((RangeQueryBuilder)builder).fieldName()
285-
: ((TermQueryBuilder)builder).fieldName();
286-
287-
List<String> incorrectTimeZones = new ArrayList<>();
288-
List<String> rewrittenFieldName = jobCaps.stream()
289-
// We only care about job caps that have the query's target field
290-
.filter(caps -> caps.getFieldCaps().keySet().contains(fieldName))
291-
.map(caps -> {
292-
RollupJobCaps.RollupFieldCaps fieldCaps = caps.getFieldCaps().get(fieldName);
293-
return fieldCaps.getAggs().stream()
294-
// For now, we only allow filtering on grouping fields
295-
.filter(agg -> {
296-
String type = (String)agg.get(RollupField.AGG);
297-
298-
// If the cap is for a date_histo, and the query is a range, the timezones need to match
299-
if (type.equals(DateHistogramAggregationBuilder.NAME) && builder instanceof RangeQueryBuilder) {
300-
String timeZone = ((RangeQueryBuilder)builder).timeZone();
301-
302-
// Many range queries don't include the timezone because the default is UTC, but the query
303-
// builder will return null so we need to set it here
304-
if (timeZone == null) {
305-
timeZone = DateTimeZone.UTC.toString();
306-
}
307-
boolean matchingTZ = ((String)agg.get(DateHistoGroupConfig.TIME_ZONE.getPreferredName()))
308-
.equalsIgnoreCase(timeZone);
309-
if (matchingTZ == false) {
310-
incorrectTimeZones.add((String)agg.get(DateHistoGroupConfig.TIME_ZONE.getPreferredName()));
311-
}
312-
return matchingTZ;
313-
}
314-
// Otherwise just make sure it's one of the three groups
315-
return type.equals(TermsAggregationBuilder.NAME)
316-
|| type.equals(DateHistogramAggregationBuilder.NAME)
317-
|| type.equals(HistogramAggregationBuilder.NAME);
318-
})
319-
// Rewrite the field name to our convention (e.g. "foo" -> "date_histogram.foo.timestamp")
320-
.map(agg -> {
321-
if (agg.get(RollupField.AGG).equals(DateHistogramAggregationBuilder.NAME)) {
322-
return RollupField.formatFieldName(fieldName, (String)agg.get(RollupField.AGG), RollupField.TIMESTAMP);
323-
} else {
324-
return RollupField.formatFieldName(fieldName, (String)agg.get(RollupField.AGG), RollupField.VALUE);
325-
}
326-
})
327-
.collect(Collectors.toList());
328-
})
329-
.distinct()
330-
.collect(ArrayList::new, List::addAll, List::addAll);
331-
332-
if (rewrittenFieldName.isEmpty()) {
333-
if (incorrectTimeZones.isEmpty()) {
334-
throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builder.getWriteableName()
335-
+ "] query is not available in selected rollup indices, cannot query.");
336-
} else {
337-
throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builder.getWriteableName()
338-
+ "] query was found in rollup indices, but requested timezone is not compatible. Options include: "
339-
+ incorrectTimeZones);
340-
}
283+
} else if (builder.getWriteableName().equals(RangeQueryBuilder.NAME)) {
284+
RangeQueryBuilder range = (RangeQueryBuilder) builder;
285+
String fieldName = range.fieldName();
286+
// Many range queries don't include the timezone because the default is UTC, but the query
287+
// builder will return null so we need to set it here
288+
String timeZone = range.timeZone() == null ? DateTimeZone.UTC.toString() : range.timeZone();
289+
290+
String rewrittenFieldName = rewriteFieldName(jobCaps, RangeQueryBuilder.NAME, fieldName, timeZone);
291+
RangeQueryBuilder rewritten = new RangeQueryBuilder(rewrittenFieldName)
292+
.from(range.from())
293+
.to(range.to())
294+
.includeLower(range.includeLower())
295+
.includeUpper(range.includeUpper());
296+
if (range.timeZone() != null) {
297+
rewritten.timeZone(range.timeZone());
341298
}
342-
343-
if (rewrittenFieldName.size() > 1) {
344-
throw new IllegalArgumentException("Ambiguous field name resolution when mapping to rolled fields. Field name [" +
345-
fieldName + "] was mapped to: [" + Strings.collectionToDelimitedString(rewrittenFieldName, ",") + "].");
299+
if (range.format() != null) {
300+
rewritten.format(range.format());
346301
}
347-
348-
//Note: instanceof here to make casting checks happier
349-
if (builder instanceof RangeQueryBuilder) {
350-
RangeQueryBuilder rewritten = new RangeQueryBuilder(rewrittenFieldName.get(0));
351-
RangeQueryBuilder original = (RangeQueryBuilder)builder;
352-
rewritten.from(original.from());
353-
rewritten.to(original.to());
354-
if (original.timeZone() != null) {
355-
rewritten.timeZone(original.timeZone());
356-
}
357-
rewritten.includeLower(original.includeLower());
358-
rewritten.includeUpper(original.includeUpper());
359-
return rewritten;
360-
} else {
361-
return new TermQueryBuilder(rewrittenFieldName.get(0), ((TermQueryBuilder)builder).value());
362-
}
363-
302+
return rewritten;
303+
} else if (builder.getWriteableName().equals(TermQueryBuilder.NAME)) {
304+
TermQueryBuilder term = (TermQueryBuilder) builder;
305+
String fieldName = term.fieldName();
306+
String rewrittenFieldName = rewriteFieldName(jobCaps, TermQueryBuilder.NAME, fieldName, null);
307+
return new TermQueryBuilder(rewrittenFieldName, term.value());
308+
} else if (builder.getWriteableName().equals(TermsQueryBuilder.NAME)) {
309+
TermsQueryBuilder terms = (TermsQueryBuilder) builder;
310+
String fieldName = terms.fieldName();
311+
String rewrittenFieldName = rewriteFieldName(jobCaps, TermQueryBuilder.NAME, fieldName, null);
312+
return new TermsQueryBuilder(rewrittenFieldName, terms.values());
364313
} else if (builder.getWriteableName().equals(MatchAllQueryBuilder.NAME)) {
365314
// no-op
366315
return builder;
@@ -369,6 +318,64 @@ static QueryBuilder rewriteQuery(QueryBuilder builder, Set<RollupJobCaps> jobCap
369318
}
370319
}
371320

321+
private static String rewriteFieldName(Set<RollupJobCaps> jobCaps,
322+
String builderName,
323+
String fieldName,
324+
String timeZone) {
325+
List<String> incompatibleTimeZones = timeZone == null ? Collections.emptyList() : new ArrayList<>();
326+
List<String> rewrittenFieldNames = jobCaps.stream()
327+
// We only care about job caps that have the query's target field
328+
.filter(caps -> caps.getFieldCaps().keySet().contains(fieldName))
329+
.map(caps -> {
330+
RollupJobCaps.RollupFieldCaps fieldCaps = caps.getFieldCaps().get(fieldName);
331+
return fieldCaps.getAggs().stream()
332+
// For now, we only allow filtering on grouping fields
333+
.filter(agg -> {
334+
String type = (String)agg.get(RollupField.AGG);
335+
336+
// If the cap is for a date_histo, and the query is a range, the timezones need to match
337+
if (type.equals(DateHistogramAggregationBuilder.NAME) && timeZone != null) {
338+
boolean matchingTZ = ((String)agg.get(DateHistoGroupConfig.TIME_ZONE.getPreferredName()))
339+
.equalsIgnoreCase(timeZone);
340+
if (matchingTZ == false) {
341+
incompatibleTimeZones.add((String)agg.get(DateHistoGroupConfig.TIME_ZONE.getPreferredName()));
342+
}
343+
return matchingTZ;
344+
}
345+
// Otherwise just make sure it's one of the three groups
346+
return type.equals(TermsAggregationBuilder.NAME)
347+
|| type.equals(DateHistogramAggregationBuilder.NAME)
348+
|| type.equals(HistogramAggregationBuilder.NAME);
349+
})
350+
// Rewrite the field name to our convention (e.g. "foo" -> "date_histogram.foo.timestamp")
351+
.map(agg -> {
352+
if (agg.get(RollupField.AGG).equals(DateHistogramAggregationBuilder.NAME)) {
353+
return RollupField.formatFieldName(fieldName, (String)agg.get(RollupField.AGG), RollupField.TIMESTAMP);
354+
} else {
355+
return RollupField.formatFieldName(fieldName, (String)agg.get(RollupField.AGG), RollupField.VALUE);
356+
}
357+
})
358+
.collect(Collectors.toList());
359+
})
360+
.distinct()
361+
.collect(ArrayList::new, List::addAll, List::addAll);
362+
if (rewrittenFieldNames.isEmpty()) {
363+
if (incompatibleTimeZones.isEmpty()) {
364+
throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builderName
365+
+ "] query is not available in selected rollup indices, cannot query.");
366+
} else {
367+
throw new IllegalArgumentException("Field [" + fieldName + "] in [" + builderName
368+
+ "] query was found in rollup indices, but requested timezone is not compatible. Options include: "
369+
+ incompatibleTimeZones);
370+
}
371+
} else if (rewrittenFieldNames.size() > 1) {
372+
throw new IllegalArgumentException("Ambiguous field name resolution when mapping to rolled fields. Field name [" +
373+
fieldName + "] was mapped to: [" + Strings.collectionToDelimitedString(rewrittenFieldNames, ",") + "].");
374+
} else {
375+
return rewrittenFieldNames.get(0);
376+
}
377+
}
378+
372379
static RollupSearchContext separateIndices(String[] indices, ImmutableOpenMap<String, IndexMetaData> indexMetaData) {
373380

374381
if (indices.length == 0) {

‎x-pack/plugin/rollup/src/test/java/org/elasticsearch/xpack/rollup/action/SearchActionTests.java

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,11 @@
2525
import org.elasticsearch.index.query.DisMaxQueryBuilder;
2626
import org.elasticsearch.index.query.MatchAllQueryBuilder;
2727
import org.elasticsearch.index.query.MatchPhraseQueryBuilder;
28+
import org.elasticsearch.index.query.MatchQueryBuilder;
2829
import org.elasticsearch.index.query.QueryBuilder;
2930
import org.elasticsearch.index.query.RangeQueryBuilder;
3031
import org.elasticsearch.index.query.TermQueryBuilder;
32+
import org.elasticsearch.index.query.TermsQueryBuilder;
3133
import org.elasticsearch.indices.IndicesModule;
3234
import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
3335
import org.elasticsearch.script.ScriptService;
@@ -61,6 +63,7 @@
6163

6264
import java.io.IOException;
6365
import java.util.ArrayList;
66+
import java.util.Arrays;
6467
import java.util.Collections;
6568
import java.util.HashMap;
6669
import java.util.HashSet;
@@ -153,7 +156,7 @@ public void testRangeWrongTZ() {
153156
"compatible. Options include: [UTC]"));
154157
}
155158

156-
public void testTerms() {
159+
public void testTermQuery() {
157160
RollupJobConfig.Builder job = ConfigTestHelpers.getRollupJob("foo");
158161
GroupConfig.Builder group = ConfigTestHelpers.getGroupConfig();
159162
group.setTerms(ConfigTestHelpers.getTerms().setFields(Collections.singletonList("foo")).build());
@@ -166,6 +169,23 @@ public void testTerms() {
166169
assertThat(((TermQueryBuilder)rewritten).fieldName(), equalTo("foo.terms.value"));
167170
}
168171

172+
public void testTermsQuery() {
173+
RollupJobConfig.Builder job = ConfigTestHelpers.getRollupJob("foo");
174+
GroupConfig.Builder group = ConfigTestHelpers.getGroupConfig();
175+
group.setTerms(ConfigTestHelpers.getTerms().setFields(Collections.singletonList("foo")).build());
176+
job.setGroupConfig(group.build());
177+
RollupJobCaps cap = new RollupJobCaps(job.build());
178+
Set<RollupJobCaps> caps = new HashSet<>();
179+
caps.add(cap);
180+
QueryBuilder original = new TermsQueryBuilder("foo", Arrays.asList("bar", "baz"));
181+
QueryBuilder rewritten =
182+
TransportRollupSearchAction.rewriteQuery(original, caps);
183+
assertThat(rewritten, instanceOf(TermsQueryBuilder.class));
184+
assertNotSame(rewritten, original);
185+
assertThat(((TermsQueryBuilder)rewritten).fieldName(), equalTo("foo.terms.value"));
186+
assertThat(((TermsQueryBuilder)rewritten).values(), equalTo(Arrays.asList("bar", "baz")));
187+
}
188+
169189
public void testCompounds() {
170190
RollupJobConfig.Builder job = ConfigTestHelpers.getRollupJob("foo");
171191
GroupConfig.Builder group = ConfigTestHelpers.getGroupConfig();

0 commit comments

Comments
 (0)