Skip to content

Commit

Permalink
Add support for numeric range keys (elastic#56452)
Browse files Browse the repository at this point in the history
This adds support for parsing numbers as range keys. They get converted
into a string, but we allow numbers.

While I was there I replaced the parser for `Range` with a
`ConstructingObjectParser` which will automatically add support for "did
you mean" style corrections on errors.

Closes elastic#56402
  • Loading branch information
nik9000 committed May 11, 2020
1 parent c85a363 commit a93dd9d
Show file tree
Hide file tree
Showing 4 changed files with 172 additions and 66 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@

import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.aggregations.AggregationBuilder;
Expand Down Expand Up @@ -50,12 +49,7 @@ public class DateRangeAggregationBuilder extends AbstractRangeBuilder<DateRangeA
for (RangeAggregator.Range range : ranges) {
agg.addRange(range);
}
}, (p, c) -> DateRangeAggregationBuilder.parseRange(p), RangeAggregator.RANGES_FIELD);
}


private static RangeAggregator.Range parseRange(XContentParser parser) throws IOException {
return RangeAggregator.Range.fromXContent(parser);
}, (p, c) -> RangeAggregator.Range.PARSER.parse(p, null), RangeAggregator.RANGES_FIELD);
}

public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@

import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.xcontent.ObjectParser;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.aggregations.AggregationBuilder;
Expand All @@ -48,11 +47,7 @@ public class RangeAggregationBuilder extends AbstractRangeBuilder<RangeAggregati
for (Range range : ranges) {
agg.addRange(range);
}
}, (p, c) -> RangeAggregationBuilder.parseRange(p), RangeAggregator.RANGES_FIELD);
}

private static Range parseRange(XContentParser parser) throws IOException {
return Range.fromXContent(parser);
}, (p, c) -> RangeAggregator.Range.PARSER.parse(p, null), RangeAggregator.RANGES_FIELD);
}

public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,12 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
import org.elasticsearch.common.xcontent.ContextParser;
import org.elasticsearch.common.xcontent.ObjectParser.ValueType;
import org.elasticsearch.common.xcontent.ToXContentObject;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.common.xcontent.XContentParserUtils;
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.search.aggregations.Aggregator;
Expand All @@ -47,6 +49,8 @@
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;

public class RangeAggregator extends BucketsAggregator {

public static final ParseField RANGES_FIELD = new ParseField("ranges");
Expand All @@ -63,6 +67,23 @@ public static class Range implements Writeable, ToXContentObject {
protected final double to;
protected final String toAsStr;

/**
* Build the range. Generally callers should prefer
* {@link Range#Range(String, Double, Double)} or
* {@link Range#Range(String, String, String)}. If you
* <strong>must</strong> call this know that consumers prefer
* {@code from} and {@code to} parameters if they are non-null
* and finite. Otherwise they parse from {@code fromrStr} and
* {@code toStr}.
*/
public Range(String key, Double from, String fromAsStr, Double to, String toAsStr) {
this.key = key;
this.from = from == null ? Double.NEGATIVE_INFINITY : from;
this.fromAsStr = fromAsStr;
this.to = to == null ? Double.POSITIVE_INFINITY : to;
this.toAsStr = toAsStr;
}

public Range(String key, Double from, Double to) {
this(key, from, null, to, null);
}
Expand Down Expand Up @@ -111,14 +132,6 @@ public String getKey() {
return this.key;
}

public Range(String key, Double from, String fromAsStr, Double to, String toAsStr) {
this.key = key;
this.from = from == null ? Double.NEGATIVE_INFINITY : from;
this.fromAsStr = fromAsStr;
this.to = to == null ? Double.POSITIVE_INFINITY : to;
this.toAsStr = toAsStr;
}

boolean matches(double value) {
return value >= from && value < to;
}
Expand All @@ -128,50 +141,6 @@ public String toString() {
return "[" + from + " to " + to + ")";
}

public static Range fromXContent(XContentParser parser) throws IOException {
XContentParser.Token token;
String currentFieldName = null;
double from = Double.NEGATIVE_INFINITY;
String fromAsStr = null;
double to = Double.POSITIVE_INFINITY;
String toAsStr = null;
String key = null;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentFieldName = parser.currentName();
} else if (token == XContentParser.Token.VALUE_NUMBER) {
if (FROM_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
from = parser.doubleValue();
} else if (TO_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
to = parser.doubleValue();
} else {
XContentParserUtils.throwUnknownField(currentFieldName, parser.getTokenLocation());
}
} else if (token == XContentParser.Token.VALUE_STRING) {
if (FROM_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
fromAsStr = parser.text();
} else if (TO_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
toAsStr = parser.text();
} else if (KEY_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
key = parser.text();
} else {
XContentParserUtils.throwUnknownField(currentFieldName, parser.getTokenLocation());
}
} else if (token == XContentParser.Token.VALUE_NULL) {
if (FROM_FIELD.match(currentFieldName, parser.getDeprecationHandler())
|| TO_FIELD.match(currentFieldName, parser.getDeprecationHandler())
|| KEY_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
// ignore null value
} else {
XContentParserUtils.throwUnknownField(currentFieldName, parser.getTokenLocation());
}
} else {
XContentParserUtils.throwUnknownToken(token, parser.getTokenLocation());
}
}
return new Range(key, from, fromAsStr, to, toAsStr);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject();
Expand All @@ -194,6 +163,35 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
return builder;
}

public static final ConstructingObjectParser<Range, Void> PARSER = new ConstructingObjectParser<>("range", arg -> {
String key = (String) arg[0];
Object from = arg[1];
Object to = arg[2];
Double fromDouble = from instanceof Number ? ((Number) from).doubleValue() : null;
Double toDouble = to instanceof Number ? ((Number) to).doubleValue() : null;
String fromStr = from instanceof String ? (String) from : null;
String toStr = to instanceof String ? (String) to : null;
return new Range(key, fromDouble, fromStr, toDouble, toStr);
});

static {
PARSER.declareField(optionalConstructorArg(),
(p, c) -> p.text(),
KEY_FIELD, ValueType.DOUBLE); // DOUBLE supports string and number
ContextParser<Void, Object> fromToParser = (p, c) -> {
if (p.currentToken() == XContentParser.Token.VALUE_STRING) {
return p.text();
}
if (p.currentToken() == XContentParser.Token.VALUE_NUMBER) {
return p.doubleValue();
}
return null;
};
// DOUBLE_OR_NULL accepts String, Number, and null
PARSER.declareField(optionalConstructorArg(), fromToParser, FROM_FIELD, ValueType.DOUBLE_OR_NULL);
PARSER.declareField(optionalConstructorArg(), fromToParser, TO_FIELD, ValueType.DOUBLE_OR_NULL);
}

@Override
public int hashCode() {
return Objects.hash(key, from, fromAsStr, to, toAsStr);
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/

package org.elasticsearch.search.aggregations.bucket.range;

import org.elasticsearch.common.io.stream.Writeable.Reader;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.common.xcontent.json.JsonXContent;
import org.elasticsearch.test.AbstractSerializingTestCase;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import static org.hamcrest.Matchers.equalTo;

public class RangeAggregationBuilderTests extends AbstractSerializingTestCase<RangeAggregationBuilder> {
@Override
protected RangeAggregationBuilder doParseInstance(XContentParser parser) throws IOException {
assertThat(parser.nextToken(), equalTo(XContentParser.Token.START_OBJECT));
assertThat(parser.nextToken(), equalTo(XContentParser.Token.FIELD_NAME));
String name = parser.currentName();
assertThat(parser.nextToken(), equalTo(XContentParser.Token.START_OBJECT));
assertThat(parser.nextToken(), equalTo(XContentParser.Token.FIELD_NAME));
assertThat(parser.currentName(), equalTo("range"));
RangeAggregationBuilder parsed = RangeAggregationBuilder.PARSER.apply(parser, name);
assertThat(parser.nextToken(), equalTo(XContentParser.Token.END_OBJECT));
assertThat(parser.nextToken(), equalTo(XContentParser.Token.END_OBJECT));
return parsed;
}

@Override
protected Reader<RangeAggregationBuilder> instanceReader() {
return RangeAggregationBuilder::new;
}

@Override
protected RangeAggregationBuilder createTestInstance() {
RangeAggregationBuilder builder = new RangeAggregationBuilder(randomAlphaOfLength(5));
builder.keyed(randomBoolean());
builder.field(randomAlphaOfLength(4));
int rangeCount = between(1, 10);
double r = 0;
for (int i = 0; i < rangeCount; i++) {
switch (between(0, 2)) {
case 0:
builder.addUnboundedFrom(randomAlphaOfLength(2), r);
break;
case 1:
builder.addUnboundedTo(randomAlphaOfLength(2), r);
break;
case 2:
double from = r;
r += randomDouble(); // less than 1
double to = r;
builder.addRange(randomAlphaOfLength(2), from, to);
break;
default:
fail();
}
r += randomDouble(); // less than 1
}
return builder;
}

@Override
protected RangeAggregationBuilder mutateInstance(RangeAggregationBuilder builder) throws IOException {
String name = builder.getName();
boolean keyed = builder.keyed();
String field = builder.field();
List<RangeAggregator.Range> ranges = builder.ranges();
switch (between(0, 3)) {
case 0:
name += randomAlphaOfLength(1);
break;
case 1:
keyed = !keyed;
break;
case 2:
field += randomAlphaOfLength(1);
break;
case 3:
ranges = new ArrayList<>(ranges);
double from = ranges.get(ranges.size() - 1).from;
double to = from + randomDouble();
ranges.add(new RangeAggregator.Range(randomAlphaOfLength(2), from, to));
break;
default:
fail();
}
RangeAggregationBuilder mutant = new RangeAggregationBuilder(name).keyed(keyed).field(field);
ranges.stream().forEach(mutant::addRange);
return mutant;
}

public void testNumericKeys() throws IOException {
RangeAggregationBuilder builder = doParseInstance(createParser(JsonXContent.jsonXContent,
"{\"test\":{\"range\":{\"field\":\"f\",\"ranges\":[{\"key\":1,\"to\":0}]}}}"));
assertThat(builder.getName(), equalTo("test"));
assertThat(builder.field(), equalTo("f"));
assertThat(builder.ranges, equalTo(org.elasticsearch.common.collect.List.of(new RangeAggregator.Range("1", null, 0d))));
}
}

0 comments on commit a93dd9d

Please sign in to comment.