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,12 @@ class org.elasticsearch.common.geo.GeoPoint {
double getLon()
}

class org.elasticsearch.common.geo.GeoBoundingBox {
org.elasticsearch.common.geo.GeoPoint topLeft()
org.elasticsearch.common.geo.GeoPoint bottomRight()
}


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

class org.elasticsearch.index.fielddata.ScriptDocValues$Geometry {
int getDimensionalType()
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,56 @@ 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 }

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

- do:
search:
rest_total_hits_as_int: true
body:
script_fields:
type:
script:
source: "doc['geo_point'].getDimensionalType()"
- match: { hits.hits.0.fields.type.0: 0 }

---
"ip":
- do:
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,230 @@
/*
* 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.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.hamcrest.Matchers;
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", this::scriptLat);
scripts.put("lon", this::scriptLon);
scripts.put("height", this::scriptHeight);
scripts.put("width", this::scriptWidth);
return scripts;
}

private double scriptHeight(Map<String, Object> vars) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
if (geometry.size() == 0) {
return Double.NaN;
} else {
GeoBoundingBox boundingBox = geometry.getBoundingBox();
return boundingBox.topLeft().lat() - boundingBox.bottomRight().lat();
}
}

private double scriptWidth(Map<String, Object> vars) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
if (geometry.size() == 0) {
return Double.NaN;
} else {
GeoBoundingBox boundingBox = geometry.getBoundingBox();
return boundingBox.bottomRight().lon() - boundingBox.topLeft().lon();
}
}

private double scriptLat(Map<String, Object> vars) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
return geometry.size() == 0 ? Double.NaN : geometry.getCentroid().lat();
}

private double scriptLon(Map<String, Object> vars) {
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
return geometry.size() == 0 ? Double.NaN : geometry.getCentroid().lon();
}

private ScriptDocValues.Geometry<?> assertGeometry(Map<?, ?> doc) {
ScriptDocValues.Geometry<?> geometry = (ScriptDocValues.Geometry<?>) doc.get("location");
if (geometry.size() == 0) {
assertThat(geometry.getBoundingBox(), Matchers.nullValue());
assertThat(geometry.getCentroid(), Matchers.nullValue());
assertThat(geometry.getDimensionalType(), equalTo(-1));
} else {
assertThat(geometry.getBoundingBox(), Matchers.notNullValue());
assertThat(geometry.getCentroid(), Matchers.notNullValue());
assertThat(geometry.getDimensionalType(), equalTo(0));
}
return geometry;
}
}

@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));
}

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));
}

public void testNullPoint() throws Exception {
client().prepareIndex("test").setId("1")
.setSource(jsonBuilder().startObject()
.field("name", "TestPosition")
.nullField("location")
.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);

Map<String, DocumentField> fields = searchResponse.getHits().getHits()[0].getFields();
assertThat(fields.get("lat").getValue(), equalTo(Double.NaN));
assertThat(fields.get("lon").getValue(), equalTo(Double.NaN));
assertThat(fields.get("height").getValue(), equalTo(Double.NaN));
assertThat(fields.get("width").getValue(), equalTo(Double.NaN));
}
}
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
Loading