Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add painless script support for Geoshape field #72886

Merged
merged 10 commits into from
May 26, 2021
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
builder.field(PRECISION.getPreferredName(), precision);
}
if (geoBoundingBox != null) {
geoBoundingBox.toXContent(builder, params);
builder.startObject(GeoBoundingBox.BOUNDS_FIELD.getPreferredName());
geoBoundingBox.toXContentFragment(builder, true);
builder.endObject();
}
builder.endObject();
return builder;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,14 @@ class org.elasticsearch.common.geo.GeoPoint {
double getLon()
}

class org.elasticsearch.common.geo.GeoBoundingBox {
double top()
double bottom()
double left()
double right()
}


class org.elasticsearch.index.fielddata.ScriptDocValues$Strings {
String get(int)
String getValue()
Expand Down Expand Up @@ -148,6 +156,11 @@ class org.elasticsearch.index.fielddata.ScriptDocValues$Doubles {
double getValue()
}

class org.elasticsearch.index.fielddata.ScriptDocValues$Geometry {
org.elasticsearch.common.geo.GeoPoint getCentroid()
org.elasticsearch.common.geo.GeoBoundingBox getBoundingBox()
}

class org.elasticsearch.index.fielddata.ScriptDocValues$GeoPoints {
org.elasticsearch.common.geo.GeoPoint get(int)
org.elasticsearch.common.geo.GeoPoint getValue()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,29 @@ setup:
- match: { hits.hits.0.fields.field.0.lat: 41.1199999647215 }
- match: { hits.hits.0.fields.field.0.lon: -71.34000004269183 }

- do:
search:
rest_total_hits_as_int: true
body:
script_fields:
centroid:
script:
source: "doc['geo_point'].getCentroid()"
- match: { hits.hits.0.fields.centroid.0.lat: 41.1199999647215 }
- match: { hits.hits.0.fields.centroid.0.lon: -71.34000004269183 }

- do:
search:
rest_total_hits_as_int: true
body:
script_fields:
bbox:
script:
source: "doc['geo_point'].getBoundingBox()"
- match: { hits.hits.0.fields.bbox.0.top_left.lat: 41.1199999647215 }
- match: { hits.hits.0.fields.bbox.0.top_left.lon: -71.34000004269183 }
- match: { hits.hits.0.fields.bbox.0.bottom_right.lat: 41.1199999647215 }
- match: { hits.hits.0.fields.bbox.0.bottom_right.lon: -71.34000004269183 }
---
"ip":
- do:
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,185 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0 and the Server Side Public License, v 1; you may not use this file except
* in compliance with, at your election, the Elastic License 2.0 or the Server
* Side Public License, v 1.
*/

package org.elasticsearch.search.geo;

import org.apache.lucene.geo.GeoEncodingUtils;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.common.document.DocumentField;
import org.elasticsearch.common.geo.GeoBoundingBox;
import org.elasticsearch.common.geo.GeoPoint;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.geo.GeometryTestUtils;
import org.elasticsearch.index.fielddata.ScriptDocValues;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.script.MockScriptPlugin;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.test.ESSingleNodeTestCase;
import org.junit.Before;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchResponse;
import static org.hamcrest.Matchers.equalTo;

public class GeoPointScriptDocValuesIT extends ESSingleNodeTestCase {

@Override
protected Collection<Class<? extends Plugin>> getPlugins() {
return Arrays.asList(CustomScriptPlugin.class);
}

public static class CustomScriptPlugin extends MockScriptPlugin {

@Override
protected Map<String, Function<Map<String, Object>, Object>> pluginScripts() {
Map<String, Function<Map<String, Object>, Object>> scripts = new HashMap<>();

scripts.put("lat", vars -> scriptLat(vars, ScriptDocValues.Geometry::getCentroid));
scripts.put("lon", vars -> scriptLon(vars, ScriptDocValues.Geometry::getCentroid));
scripts.put("height", vars -> scriptHeight(vars, ScriptDocValues.Geometry::getBoundingBox));
scripts.put("width", vars -> scriptWidth(vars, ScriptDocValues.Geometry::getBoundingBox));
return scripts;
}

static Double scriptHeight(Map<String, Object> vars, Function<ScriptDocValues.Geometry<?>, GeoBoundingBox> bbox) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
GeoBoundingBox boundingBox = bbox.apply((ScriptDocValues.Geometry<?>) doc.get("location"));
return boundingBox.top() - boundingBox.bottom();
}

static Double scriptWidth(Map<String, Object> vars, Function<ScriptDocValues.Geometry<?>, GeoBoundingBox> bbox) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
GeoBoundingBox boundingBox = bbox.apply((ScriptDocValues.Geometry<?>) doc.get("location"));
return boundingBox.right() - boundingBox.left();
}

static Double scriptLat(Map<String, Object> vars, Function<ScriptDocValues.Geometry<?>, GeoPoint> centroid) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
return centroid.apply((ScriptDocValues.Geometry<?>) doc.get("location")).lat();
}

static Double scriptLon(Map<String, Object> vars, Function<ScriptDocValues.Geometry<?>, GeoPoint> centroid) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
return centroid.apply((ScriptDocValues.Geometry<?>) doc.get("location")).lon();
}
}

@Override
protected boolean forbidPrivateIndexSettings() {
return false;
}

@Before
public void setupTestIndex() throws IOException {
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder().startObject().startObject("_doc")
.startObject("properties").startObject("location").field("type", "geo_point");
xContentBuilder.endObject().endObject().endObject().endObject();
assertAcked(client().admin().indices().prepareCreate("test").setMapping(xContentBuilder));
ensureGreen();
}

public void testRandomPoint() throws Exception {
final double lat = GeometryTestUtils.randomLat();
final double lon = GeometryTestUtils.randomLon();
client().prepareIndex("test").setId("1")
.setSource(jsonBuilder().startObject()
.field("name", "TestPosition")
.field("location", new double[]{lon, lat})
.endObject())
.get();

client().admin().indices().prepareRefresh("test").get();

SearchResponse searchResponse = client().prepareSearch().addStoredField("_source")
.addScriptField("lat", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lat", Collections.emptyMap()))
.addScriptField("lon", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lon", Collections.emptyMap()))
.addScriptField("height", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "height", Collections.emptyMap()))
.addScriptField("width", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "width", Collections.emptyMap()))
.get();
assertSearchResponse(searchResponse);

final double qLat = GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(lat));
final double qLon = GeoEncodingUtils.decodeLongitude(GeoEncodingUtils.encodeLongitude(lon));

Map<String, DocumentField> fields = searchResponse.getHits().getHits()[0].getFields();
assertThat(fields.get("lat").getValue(), equalTo(qLat));
assertThat(fields.get("lon").getValue(), equalTo(qLon));
assertThat(fields.get("height").getValue(), equalTo(0d));
assertThat(fields.get("width").getValue(), equalTo(0d));

client().prepareDelete("test", "1").get();
client().admin().indices().prepareRefresh("test").get();
}

public void testRandomMultiPoint() throws Exception {
final int size = randomIntBetween(2, 20);
final double[] lats = new double[size];
final double[] lons = new double[size];
for (int i = 0; i < size; i++) {
lats[i] = GeometryTestUtils.randomLat();
lons[i] = GeometryTestUtils.randomLon();
}

final double[][] values = new double[size][];
for (int i = 0; i < size; i++) {
values[i] = new double[]{lons[i], lats[i]};
}

XContentBuilder builder = jsonBuilder().startObject()
.field("name", "TestPosition")
.field("location", values).endObject();
client().prepareIndex("test").setId("1").setSource(builder).get();

client().admin().indices().prepareRefresh("test").get();

SearchResponse searchResponse = client().prepareSearch().addStoredField("_source")
.addScriptField("lat", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lat", Collections.emptyMap()))
.addScriptField("lon", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lon", Collections.emptyMap()))
.addScriptField("height", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "height", Collections.emptyMap()))
.addScriptField("width", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "width", Collections.emptyMap()))
.get();
assertSearchResponse(searchResponse);

for (int i = 0; i < size; i++) {
lats[i] = GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(lats[i]));
lons[i] = GeoEncodingUtils.decodeLongitude(GeoEncodingUtils.encodeLongitude(lons[i]));
}

final double centroidLon = Arrays.stream(lons).sum() / size;
final double centroidLat = Arrays.stream(lats).sum() / size;
final double width = Arrays.stream(lons).max().getAsDouble() - Arrays.stream(lons).min().getAsDouble();
final double height = Arrays.stream(lats).max().getAsDouble() - Arrays.stream(lats).min().getAsDouble();

Map<String, DocumentField> fields = searchResponse.getHits().getHits()[0].getFields();
assertThat(fields.get("lat").getValue(), equalTo(centroidLat));
assertThat(fields.get("lon").getValue(), equalTo(centroidLon));
assertThat(fields.get("height").getValue(), equalTo(height));
assertThat(fields.get("width").getValue(), equalTo(width));

client().prepareDelete("test", "1").get();
client().admin().indices().prepareRefresh("test").get();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
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.ToXContentObject;
import org.elasticsearch.common.xcontent.ToXContentFragment;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.geometry.Geometry;
Expand All @@ -29,7 +29,7 @@
* A class representing a Geo-Bounding-Box for use by Geo queries and aggregations
* that deal with extents/rectangles representing rectangular areas of interest.
*/
public class GeoBoundingBox implements ToXContentObject, Writeable {
public class GeoBoundingBox implements ToXContentFragment, Writeable {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have to make GeoBoundingBox to behave more like a GeoPoint

private static final WellKnownText WKT_PARSER = new WellKnownText(true, new StandardValidator(true));
static final ParseField TOP_RIGHT_FIELD = new ParseField("top_right");
static final ParseField BOTTOM_LEFT_FIELD = new ParseField("bottom_left");
Expand Down Expand Up @@ -88,7 +88,7 @@ public double right() {

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject(BOUNDS_FIELD.getPreferredName());
builder.startObject();
toXContentFragment(builder, true);
builder.endObject();
return builder;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefBuilder;
import org.elasticsearch.common.geo.GeoBoundingBox;
import org.elasticsearch.common.geo.GeoPoint;
import org.elasticsearch.common.geo.GeoUtils;
import org.elasticsearch.common.time.DateUtils;
Expand Down Expand Up @@ -247,10 +248,20 @@ public int size() {
}
}

public static final class GeoPoints extends ScriptDocValues<GeoPoint> {
public abstract static class Geometry<T> extends ScriptDocValues<T> {
/** Returns the centroid of this geometry */
public abstract GeoPoint getCentroid();
/** Returns the bounding box of this geometry */
public abstract GeoBoundingBox getBoundingBox();

}

public static final class GeoPoints extends Geometry<GeoPoint> {

private final MultiGeoPointValues in;
private GeoPoint[] values = new GeoPoint[0];
private final GeoPoint centroid = new GeoPoint();
private final GeoBoundingBox boundingBox = new GeoBoundingBox(new GeoPoint(), new GeoPoint());
private int count;

public GeoPoints(MultiGeoPointValues in) {
Expand All @@ -261,10 +272,25 @@ public GeoPoints(MultiGeoPointValues in) {
public void setNextDocId(int docId) throws IOException {
if (in.advanceExact(docId)) {
resize(in.docValueCount());
double centroidLat = 0;
double centroidLon = 0;
double maxLon = Double.NEGATIVE_INFINITY;
double minLon = Double.POSITIVE_INFINITY;
double maxLat = Double.NEGATIVE_INFINITY;
double minLat = Double.POSITIVE_INFINITY;
for (int i = 0; i < count; i++) {
GeoPoint point = in.nextValue();
values[i] = new GeoPoint(point.lat(), point.lon());
values[i].reset(point.lat(), point.lon());
centroidLat += point.getLat();
centroidLon += point.getLon();
maxLon = Math.max(maxLon, values[i].getLon());
minLon = Math.min(minLon, values[i].getLon());
maxLat = Math.max(maxLat, values[i].getLat());
minLat = Math.min(minLat, values[i].getLat());
}
centroid.reset(centroidLat / count, centroidLon / count);
boundingBox.topLeft().reset(maxLat, minLon);
boundingBox.bottomRight().reset(minLat, maxLon);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One of my worries here is to make the construction of a GeoPoint script value expensive.

On the other hand it seems there is a small bug here as we were building new GeoPoint objects for each new doc value which seems wasteful.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given this is an internal detail we could always attempt to improve this at a later time.

} else {
resize(0);
}
Expand All @@ -280,7 +306,7 @@ protected void resize(int newSize) {
int oldLength = values.length;
values = ArrayUtil.grow(values, count);
for (int i = oldLength; i < values.length; ++i) {
values[i] = new GeoPoint();
values[i] = new GeoPoint();
}
}
}
Expand Down Expand Up @@ -364,6 +390,16 @@ public double geohashDistanceWithDefault(String geohash, double defaultValue) {
}
return geohashDistance(geohash);
}

@Override
public GeoPoint getCentroid() {
return centroid;
}

@Override
public GeoBoundingBox getBoundingBox() {
return boundingBox;
}
}

public static final class Booleans extends ScriptDocValues<Boolean> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,9 @@ protected void innerWriteTo(StreamOutput out) throws IOException {
protected void doXContentBody(XContentBuilder builder, Params params) throws IOException {
builder.field("precision", precision);
if (geoBoundingBox.isUnbounded() == false) {
geoBoundingBox.toXContent(builder, params);
builder.startObject(GeoBoundingBox.BOUNDS_FIELD.getPreferredName());
geoBoundingBox.toXContentFragment(builder, true);
builder.endObject();
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -201,7 +201,9 @@ protected XContentBuilder doXContentBody(XContentBuilder builder, Params params)
builder.field(FIELD_SHARD_SIZE.getPreferredName(), shardSize);
}
if (geoBoundingBox.isUnbounded() == false) {
geoBoundingBox.toXContent(builder, params);
builder.startObject(GeoBoundingBox.BOUNDS_FIELD.getPreferredName());
geoBoundingBox.toXContentFragment(builder, true);
builder.endObject();
}
return builder;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,9 @@ public Object getProperty(List<String> path) {
public XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
GeoBoundingBox bbox = resolveGeoBoundingBox();
if (bbox != null) {
bbox.toXContent(builder, params);
builder.startObject(GeoBoundingBox.BOUNDS_FIELD.getPreferredName());
bbox.toXContentFragment(builder, true);
builder.endObject();
}
return builder;
}
Expand Down
Loading