Skip to content

Commit e4f45db

Browse files
author
Hendrik Muhs
authored
[Transform] add support for script in group_by (#53167)
add the possibility to base the group_by on the output of a script. closes #43152
1 parent 54b66d3 commit e4f45db

File tree

25 files changed

+847
-283
lines changed

25 files changed

+847
-283
lines changed

‎client/rest-high-level/src/main/java/org/elasticsearch/client/transform/transforms/pivot/DateHistogramGroupSource.java

Lines changed: 77 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,14 @@
2020
package org.elasticsearch.client.transform.transforms.pivot;
2121

2222
import org.elasticsearch.common.ParseField;
23+
import org.elasticsearch.common.Strings;
2324
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
2425
import org.elasticsearch.common.xcontent.ObjectParser;
2526
import org.elasticsearch.common.xcontent.ToXContentFragment;
2627
import org.elasticsearch.common.xcontent.ToXContentObject;
2728
import org.elasticsearch.common.xcontent.XContentBuilder;
2829
import org.elasticsearch.common.xcontent.XContentParser;
30+
import org.elasticsearch.script.Script;
2931
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;
3032

3133
import java.io.IOException;
@@ -48,23 +50,28 @@ public class DateHistogramGroupSource extends SingleGroupSource implements ToXCo
4850

4951
// From DateHistogramAggregationBuilder in core, transplanted and modified to a set
5052
// so we don't need to import a dependency on the class
51-
private static final Set<String> DATE_FIELD_UNITS = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(
52-
"year",
53-
"1y",
54-
"quarter",
55-
"1q",
56-
"month",
57-
"1M",
58-
"week",
59-
"1w",
60-
"day",
61-
"1d",
62-
"hour",
63-
"1h",
64-
"minute",
65-
"1m",
66-
"second",
67-
"1s")));
53+
private static final Set<String> DATE_FIELD_UNITS = Collections.unmodifiableSet(
54+
new HashSet<>(
55+
Arrays.asList(
56+
"year",
57+
"1y",
58+
"quarter",
59+
"1q",
60+
"month",
61+
"1M",
62+
"week",
63+
"1w",
64+
"day",
65+
"1d",
66+
"hour",
67+
"1h",
68+
"minute",
69+
"1m",
70+
"second",
71+
"1s"
72+
)
73+
)
74+
);
6875

6976
/**
7077
* Interval can be specified in 2 ways:
@@ -76,6 +83,7 @@ public class DateHistogramGroupSource extends SingleGroupSource implements ToXCo
7683
*/
7784
public interface Interval extends ToXContentFragment {
7885
String getName();
86+
7987
DateHistogramInterval getInterval();
8088
}
8189

@@ -131,8 +139,9 @@ public static class CalendarInterval implements Interval {
131139
public CalendarInterval(DateHistogramInterval interval) {
132140
this.interval = interval;
133141
if (DATE_FIELD_UNITS.contains(interval.toString()) == false) {
134-
throw new IllegalArgumentException("The supplied interval [" + interval + "] could not be parsed " +
135-
"as a calendar interval.");
142+
throw new IllegalArgumentException(
143+
"The supplied interval [" + interval + "] could not be parsed " + "as a calendar interval."
144+
);
136145
}
137146
}
138147

@@ -173,33 +182,35 @@ public int hashCode() {
173182
}
174183
}
175184

176-
private static final ConstructingObjectParser<DateHistogramGroupSource, Void> PARSER =
177-
new ConstructingObjectParser<>("date_histogram_group_source",
178-
true,
179-
(args) -> {
180-
String field = (String)args[0];
181-
String fixedInterval = (String) args[1];
182-
String calendarInterval = (String) args[2];
183-
184-
Interval interval = null;
185-
186-
if (fixedInterval != null && calendarInterval != null) {
187-
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found both");
188-
} else if (fixedInterval != null) {
189-
interval = new FixedInterval(new DateHistogramInterval(fixedInterval));
190-
} else if (calendarInterval != null) {
191-
interval = new CalendarInterval(new DateHistogramInterval(calendarInterval));
192-
} else {
193-
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found none");
194-
}
195-
196-
ZoneId zoneId = (ZoneId) args[3];
197-
return new DateHistogramGroupSource(field, interval, zoneId);
198-
});
185+
private static final ConstructingObjectParser<DateHistogramGroupSource, Void> PARSER = new ConstructingObjectParser<>(
186+
"date_histogram_group_source",
187+
true,
188+
(args) -> {
189+
String field = (String) args[0];
190+
Script script = (Script) args[1];
191+
String fixedInterval = (String) args[2];
192+
String calendarInterval = (String) args[3];
193+
ZoneId zoneId = (ZoneId) args[4];
194+
195+
Interval interval = null;
196+
197+
if (fixedInterval != null && calendarInterval != null) {
198+
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found both");
199+
} else if (fixedInterval != null) {
200+
interval = new FixedInterval(new DateHistogramInterval(fixedInterval));
201+
} else if (calendarInterval != null) {
202+
interval = new CalendarInterval(new DateHistogramInterval(calendarInterval));
203+
} else {
204+
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found none");
205+
}
206+
207+
return new DateHistogramGroupSource(field, script, interval, zoneId);
208+
}
209+
);
199210

200211
static {
201212
PARSER.declareString(optionalConstructorArg(), FIELD);
202-
213+
Script.declareScript(PARSER, optionalConstructorArg(), SCRIPT);
203214
PARSER.declareString(optionalConstructorArg(), new ParseField(FixedInterval.NAME));
204215
PARSER.declareString(optionalConstructorArg(), new ParseField(CalendarInterval.NAME));
205216

@@ -219,8 +230,8 @@ public static DateHistogramGroupSource fromXContent(final XContentParser parser)
219230
private final Interval interval;
220231
private final ZoneId timeZone;
221232

222-
DateHistogramGroupSource(String field, Interval interval, ZoneId timeZone) {
223-
super(field);
233+
DateHistogramGroupSource(String field, Script script, Interval interval, ZoneId timeZone) {
234+
super(field, script);
224235
this.interval = interval;
225236
this.timeZone = timeZone;
226237
}
@@ -241,9 +252,7 @@ public ZoneId getTimeZone() {
241252
@Override
242253
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
243254
builder.startObject();
244-
if (field != null) {
245-
builder.field(FIELD.getPreferredName(), field);
246-
}
255+
super.innerXContent(builder, params);
247256
interval.toXContent(builder, params);
248257
if (timeZone != null) {
249258
builder.field(TIME_ZONE.getPreferredName(), timeZone.toString());
@@ -264,23 +273,29 @@ public boolean equals(Object other) {
264273

265274
final DateHistogramGroupSource that = (DateHistogramGroupSource) other;
266275

267-
return Objects.equals(this.field, that.field) &&
268-
Objects.equals(this.interval, that.interval) &&
269-
Objects.equals(this.timeZone, that.timeZone);
276+
return Objects.equals(this.field, that.field)
277+
&& Objects.equals(this.interval, that.interval)
278+
&& Objects.equals(this.timeZone, that.timeZone);
270279
}
271280

272281
@Override
273282
public int hashCode() {
274283
return Objects.hash(field, interval, timeZone);
275284
}
276285

286+
@Override
287+
public String toString() {
288+
return Strings.toString(this, true, true);
289+
}
290+
277291
public static Builder builder() {
278292
return new Builder();
279293
}
280294

281295
public static class Builder {
282296

283297
private String field;
298+
private Script script;
284299
private Interval interval;
285300
private ZoneId timeZone;
286301

@@ -294,6 +309,16 @@ public Builder setField(String field) {
294309
return this;
295310
}
296311

312+
/**
313+
* The script with which to construct the date histogram grouping
314+
* @param script The script
315+
* @return The {@link Builder} with the script set.
316+
*/
317+
public Builder setScript(Script script) {
318+
this.script = script;
319+
return this;
320+
}
321+
297322
/**
298323
* Set the interval for the DateHistogram grouping
299324
* @param interval a fixed or calendar interval
@@ -315,7 +340,7 @@ public Builder setTimeZone(ZoneId timeZone) {
315340
}
316341

317342
public DateHistogramGroupSource build() {
318-
return new DateHistogramGroupSource(field, interval, timeZone);
343+
return new DateHistogramGroupSource(field, script, interval, timeZone);
319344
}
320345
}
321346
}

‎client/rest-high-level/src/main/java/org/elasticsearch/client/transform/transforms/pivot/HistogramGroupSource.java

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.elasticsearch.common.xcontent.ToXContentObject;
2626
import org.elasticsearch.common.xcontent.XContentBuilder;
2727
import org.elasticsearch.common.xcontent.XContentParser;
28+
import org.elasticsearch.script.Script;
2829

2930
import java.io.IOException;
3031
import java.util.Objects;
@@ -37,12 +38,15 @@
3738
public class HistogramGroupSource extends SingleGroupSource implements ToXContentObject {
3839

3940
protected static final ParseField INTERVAL = new ParseField("interval");
40-
private static final ConstructingObjectParser<HistogramGroupSource, Void> PARSER =
41-
new ConstructingObjectParser<>("histogram_group_source", true,
42-
args -> new HistogramGroupSource((String) args[0], (double) args[1]));
41+
private static final ConstructingObjectParser<HistogramGroupSource, Void> PARSER = new ConstructingObjectParser<>(
42+
"histogram_group_source",
43+
true,
44+
args -> new HistogramGroupSource((String) args[0], (Script) args[1], (double) args[2])
45+
);
4346

4447
static {
4548
PARSER.declareString(optionalConstructorArg(), FIELD);
49+
Script.declareScript(PARSER, optionalConstructorArg(), SCRIPT);
4650
PARSER.declareDouble(optionalConstructorArg(), INTERVAL);
4751
}
4852

@@ -52,8 +56,8 @@ public static HistogramGroupSource fromXContent(final XContentParser parser) {
5256

5357
private final double interval;
5458

55-
HistogramGroupSource(String field, double interval) {
56-
super(field);
59+
HistogramGroupSource(String field, Script script, double interval) {
60+
super(field, script);
5761
if (interval <= 0) {
5862
throw new IllegalArgumentException("[interval] must be greater than 0.");
5963
}
@@ -72,9 +76,7 @@ public double getInterval() {
7276
@Override
7377
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
7478
builder.startObject();
75-
if (field != null) {
76-
builder.field(FIELD.getPreferredName(), field);
77-
}
79+
super.innerXContent(builder, params);
7880
builder.field(INTERVAL.getPreferredName(), interval);
7981
builder.endObject();
8082
return builder;
@@ -92,8 +94,7 @@ public boolean equals(Object other) {
9294

9395
final HistogramGroupSource that = (HistogramGroupSource) other;
9496

95-
return Objects.equals(this.field, that.field) &&
96-
Objects.equals(this.interval, that.interval);
97+
return Objects.equals(this.field, that.field) && Objects.equals(this.interval, that.interval);
9798
}
9899

99100
@Override
@@ -108,6 +109,7 @@ public static Builder builder() {
108109
public static class Builder {
109110

110111
private String field;
112+
private Script script;
111113
private double interval;
112114

113115
/**
@@ -130,8 +132,18 @@ public Builder setInterval(double interval) {
130132
return this;
131133
}
132134

135+
/**
136+
* The script with which to construct the histogram grouping
137+
* @param script The script
138+
* @return The {@link Builder} with the script set.
139+
*/
140+
public Builder setScript(Script script) {
141+
this.script = script;
142+
return this;
143+
}
144+
133145
public HistogramGroupSource build() {
134-
return new HistogramGroupSource(field, interval);
146+
return new HistogramGroupSource(field, script, interval);
135147
}
136148
}
137149
}

‎client/rest-high-level/src/main/java/org/elasticsearch/client/transform/transforms/pivot/SingleGroupSource.java

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,17 @@
2121

2222
import org.elasticsearch.common.ParseField;
2323
import org.elasticsearch.common.xcontent.ToXContentObject;
24+
import org.elasticsearch.common.xcontent.XContentBuilder;
25+
import org.elasticsearch.script.Script;
2426

27+
import java.io.IOException;
2528
import java.util.Locale;
2629
import java.util.Objects;
2730

2831
public abstract class SingleGroupSource implements ToXContentObject {
2932

3033
protected static final ParseField FIELD = new ParseField("field");
34+
protected static final ParseField SCRIPT = new ParseField("script");
3135

3236
public enum Type {
3337
TERMS,
@@ -40,9 +44,11 @@ public String value() {
4044
}
4145

4246
protected final String field;
47+
protected final Script script;
4348

44-
public SingleGroupSource(final String field) {
49+
public SingleGroupSource(final String field, final Script script) {
4550
this.field = field;
51+
this.script = script;
4652
}
4753

4854
public abstract Type getType();
@@ -51,6 +57,19 @@ public String getField() {
5157
return field;
5258
}
5359

60+
public Script getScript() {
61+
return script;
62+
}
63+
64+
protected void innerXContent(XContentBuilder builder, Params params) throws IOException {
65+
if (field != null) {
66+
builder.field(FIELD.getPreferredName(), field);
67+
}
68+
if (script != null) {
69+
builder.field(SCRIPT.getPreferredName(), script);
70+
}
71+
}
72+
5473
@Override
5574
public boolean equals(Object other) {
5675
if (this == other) {
@@ -63,11 +82,11 @@ public boolean equals(Object other) {
6382

6483
final SingleGroupSource that = (SingleGroupSource) other;
6584

66-
return Objects.equals(this.field, that.field);
85+
return Objects.equals(this.field, that.field) && Objects.equals(this.script, that.script);
6786
}
6887

6988
@Override
7089
public int hashCode() {
71-
return Objects.hash(field);
90+
return Objects.hash(field, script);
7291
}
7392
}

0 commit comments

Comments
 (0)