Skip to content

Commit

Permalink
add support for script in group_by
Browse files Browse the repository at this point in the history
  • Loading branch information
Hendrik Muhs committed Mar 5, 2020
1 parent 8980120 commit 32a2ae7
Show file tree
Hide file tree
Showing 25 changed files with 848 additions and 283 deletions.
129 changes: 77 additions & 52 deletions ...in/java/org/elasticsearch/client/transform/transforms/pivot/DateHistogramGroupSource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,14 @@
package org.elasticsearch.client.transform.transforms.pivot;

import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.ToXContentFragment;
import org.elasticsearch.common.xcontent.ToXContentObject;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;

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

// From DateHistogramAggregationBuilder in core, transplanted and modified to a set
// so we don't need to import a dependency on the class
private static final Set<String> DATE_FIELD_UNITS = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(
"year",
"1y",
"quarter",
"1q",
"month",
"1M",
"week",
"1w",
"day",
"1d",
"hour",
"1h",
"minute",
"1m",
"second",
"1s")));
private static final Set<String> DATE_FIELD_UNITS = Collections.unmodifiableSet(
new HashSet<>(
Arrays.asList(
"year",
"1y",
"quarter",
"1q",
"month",
"1M",
"week",
"1w",
"day",
"1d",
"hour",
"1h",
"minute",
"1m",
"second",
"1s"
)
)
);

/**
* Interval can be specified in 2 ways:
Expand All @@ -76,6 +83,7 @@ public class DateHistogramGroupSource extends SingleGroupSource implements ToXCo
*/
public interface Interval extends ToXContentFragment {
String getName();

DateHistogramInterval getInterval();
}

Expand Down Expand Up @@ -131,8 +139,9 @@ public static class CalendarInterval implements Interval {
public CalendarInterval(DateHistogramInterval interval) {
this.interval = interval;
if (DATE_FIELD_UNITS.contains(interval.toString()) == false) {
throw new IllegalArgumentException("The supplied interval [" + interval + "] could not be parsed " +
"as a calendar interval.");
throw new IllegalArgumentException(
"The supplied interval [" + interval + "] could not be parsed " + "as a calendar interval."
);
}
}

Expand Down Expand Up @@ -173,33 +182,35 @@ public int hashCode() {
}
}

private static final ConstructingObjectParser<DateHistogramGroupSource, Void> PARSER =
new ConstructingObjectParser<>("date_histogram_group_source",
true,
(args) -> {
String field = (String)args[0];
String fixedInterval = (String) args[1];
String calendarInterval = (String) args[2];

Interval interval = null;

if (fixedInterval != null && calendarInterval != null) {
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found both");
} else if (fixedInterval != null) {
interval = new FixedInterval(new DateHistogramInterval(fixedInterval));
} else if (calendarInterval != null) {
interval = new CalendarInterval(new DateHistogramInterval(calendarInterval));
} else {
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found none");
}

ZoneId zoneId = (ZoneId) args[3];
return new DateHistogramGroupSource(field, interval, zoneId);
});
private static final ConstructingObjectParser<DateHistogramGroupSource, Void> PARSER = new ConstructingObjectParser<>(
"date_histogram_group_source",
true,
(args) -> {
String field = (String) args[0];
Script script = (Script) args[1];
String fixedInterval = (String) args[2];
String calendarInterval = (String) args[3];
ZoneId zoneId = (ZoneId) args[4];

Interval interval = null;

if (fixedInterval != null && calendarInterval != null) {
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found both");
} else if (fixedInterval != null) {
interval = new FixedInterval(new DateHistogramInterval(fixedInterval));
} else if (calendarInterval != null) {
interval = new CalendarInterval(new DateHistogramInterval(calendarInterval));
} else {
throw new IllegalArgumentException("You must specify either fixed_interval or calendar_interval, found none");
}

return new DateHistogramGroupSource(field, script, interval, zoneId);
}
);

static {
PARSER.declareString(optionalConstructorArg(), FIELD);

Script.declareScript(PARSER, optionalConstructorArg(), SCRIPT);
PARSER.declareString(optionalConstructorArg(), new ParseField(FixedInterval.NAME));
PARSER.declareString(optionalConstructorArg(), new ParseField(CalendarInterval.NAME));

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

DateHistogramGroupSource(String field, Interval interval, ZoneId timeZone) {
super(field);
DateHistogramGroupSource(String field, Script script, Interval interval, ZoneId timeZone) {
super(field, script);
this.interval = interval;
this.timeZone = timeZone;
}
Expand All @@ -241,9 +252,7 @@ public ZoneId getTimeZone() {
@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject();
if (field != null) {
builder.field(FIELD.getPreferredName(), field);
}
super.innerXContent(builder, params);
interval.toXContent(builder, params);
if (timeZone != null) {
builder.field(TIME_ZONE.getPreferredName(), timeZone.toString());
Expand All @@ -264,23 +273,29 @@ public boolean equals(Object other) {

final DateHistogramGroupSource that = (DateHistogramGroupSource) other;

return Objects.equals(this.field, that.field) &&
Objects.equals(this.interval, that.interval) &&
Objects.equals(this.timeZone, that.timeZone);
return Objects.equals(this.field, that.field)
&& Objects.equals(this.interval, that.interval)
&& Objects.equals(this.timeZone, that.timeZone);
}

@Override
public int hashCode() {
return Objects.hash(field, interval, timeZone);
}

@Override
public String toString() {
return Strings.toString(this, true, true);
}

public static Builder builder() {
return new Builder();
}

public static class Builder {

private String field;
private Script script;
private Interval interval;
private ZoneId timeZone;

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

/**
* The script with which to construct the date histogram grouping
* @param script The script
* @return The {@link Builder} with the script set.
*/
public Builder setScript(Script script) {
this.script = script;
return this;
}

/**
* Set the interval for the DateHistogram grouping
* @param interval a fixed or calendar interval
Expand All @@ -315,7 +340,7 @@ public Builder setTimeZone(ZoneId timeZone) {
}

public DateHistogramGroupSource build() {
return new DateHistogramGroupSource(field, interval, timeZone);
return new DateHistogramGroupSource(field, script, interval, timeZone);
}
}
}
34 changes: 23 additions & 11 deletions ...c/main/java/org/elasticsearch/client/transform/transforms/pivot/HistogramGroupSource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import org.elasticsearch.common.xcontent.ToXContentObject;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.script.Script;

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

protected static final ParseField INTERVAL = new ParseField("interval");
private static final ConstructingObjectParser<HistogramGroupSource, Void> PARSER =
new ConstructingObjectParser<>("histogram_group_source", true,
args -> new HistogramGroupSource((String) args[0], (double) args[1]));
private static final ConstructingObjectParser<HistogramGroupSource, Void> PARSER = new ConstructingObjectParser<>(
"histogram_group_source",
true,
args -> new HistogramGroupSource((String) args[0], (Script) args[1], (double) args[2])
);

static {
PARSER.declareString(optionalConstructorArg(), FIELD);
Script.declareScript(PARSER, optionalConstructorArg(), SCRIPT);
PARSER.declareDouble(optionalConstructorArg(), INTERVAL);
}

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

private final double interval;

HistogramGroupSource(String field, double interval) {
super(field);
HistogramGroupSource(String field, Script script, double interval) {
super(field, script);
if (interval <= 0) {
throw new IllegalArgumentException("[interval] must be greater than 0.");
}
Expand All @@ -72,9 +76,7 @@ public double getInterval() {
@Override
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
builder.startObject();
if (field != null) {
builder.field(FIELD.getPreferredName(), field);
}
super.innerXContent(builder, params);
builder.field(INTERVAL.getPreferredName(), interval);
builder.endObject();
return builder;
Expand All @@ -92,8 +94,7 @@ public boolean equals(Object other) {

final HistogramGroupSource that = (HistogramGroupSource) other;

return Objects.equals(this.field, that.field) &&
Objects.equals(this.interval, that.interval);
return Objects.equals(this.field, that.field) && Objects.equals(this.interval, that.interval);
}

@Override
Expand All @@ -108,6 +109,7 @@ public static Builder builder() {
public static class Builder {

private String field;
private Script script;
private double interval;

/**
Expand All @@ -130,8 +132,18 @@ public Builder setInterval(double interval) {
return this;
}

/**
* The script with which to construct the histogram grouping
* @param script The script
* @return The {@link Builder} with the script set.
*/
public Builder setScript(Script script) {
this.script = script;
return this;
}

public HistogramGroupSource build() {
return new HistogramGroupSource(field, interval);
return new HistogramGroupSource(field, script, interval);
}
}
}
25 changes: 22 additions & 3 deletions .../src/main/java/org/elasticsearch/client/transform/transforms/pivot/SingleGroupSource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,17 @@

import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.xcontent.ToXContentObject;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.script.Script;

import java.io.IOException;
import java.util.Locale;
import java.util.Objects;

public abstract class SingleGroupSource implements ToXContentObject {

protected static final ParseField FIELD = new ParseField("field");
protected static final ParseField SCRIPT = new ParseField("script");

public enum Type {
TERMS,
Expand All @@ -40,9 +44,11 @@ public String value() {
}

protected final String field;
protected final Script script;

public SingleGroupSource(final String field) {
public SingleGroupSource(final String field, final Script script) {
this.field = field;
this.script = script;
}

public abstract Type getType();
Expand All @@ -51,6 +57,19 @@ public String getField() {
return field;
}

public Script getScript() {
return script;
}

protected void innerXContent(XContentBuilder builder, Params params) throws IOException {
if (field != null) {
builder.field(FIELD.getPreferredName(), field);
}
if (script != null) {
builder.field(SCRIPT.getPreferredName(), script);
}
}

@Override
public boolean equals(Object other) {
if (this == other) {
Expand All @@ -63,11 +82,11 @@ public boolean equals(Object other) {

final SingleGroupSource that = (SingleGroupSource) other;

return Objects.equals(this.field, that.field);
return Objects.equals(this.field, that.field) && Objects.equals(this.script, that.script);
}

@Override
public int hashCode() {
return Objects.hash(field);
return Objects.hash(field, script);
}
}
Loading

0 comments on commit 32a2ae7

Please sign in to comment.