Skip to content

Commit f926d42

Browse files
iveraseiverase
authored
Add painless script support for geo_shape field (#72886) (#73432)
Users can access the centroid, bounding box and dimensional type of the shape.
1 parent 9005141 commit f926d42

File tree

17 files changed

+736
-36
lines changed

17 files changed

+736
-36
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
9090
builder.field(PRECISION.getPreferredName(), precision);
9191
}
9292
if (geoBoundingBox != null) {
93-
geoBoundingBox.toXContent(builder, params);
93+
builder.startObject(GeoBoundingBox.BOUNDS_FIELD.getPreferredName());
94+
geoBoundingBox.toXContentFragment(builder, true);
95+
builder.endObject();
9496
}
9597
builder.endObject();
9698
return builder;

‎modules/lang-painless/src/main/resources/org/elasticsearch/painless/spi/org.elasticsearch.txt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,12 @@ class org.elasticsearch.common.geo.GeoPoint {
5353
double getLon()
5454
}
5555

56+
class org.elasticsearch.common.geo.GeoBoundingBox {
57+
org.elasticsearch.common.geo.GeoPoint topLeft()
58+
org.elasticsearch.common.geo.GeoPoint bottomRight()
59+
}
60+
61+
5662
class org.elasticsearch.index.fielddata.ScriptDocValues$Strings {
5763
String get(int)
5864
String getValue()
@@ -148,6 +154,12 @@ class org.elasticsearch.index.fielddata.ScriptDocValues$Doubles {
148154
double getValue()
149155
}
150156

157+
class org.elasticsearch.index.fielddata.ScriptDocValues$Geometry {
158+
int getDimensionalType()
159+
org.elasticsearch.common.geo.GeoPoint getCentroid()
160+
org.elasticsearch.common.geo.GeoBoundingBox getBoundingBox()
161+
}
162+
151163
class org.elasticsearch.index.fielddata.ScriptDocValues$GeoPoints {
152164
org.elasticsearch.common.geo.GeoPoint get(int)
153165
org.elasticsearch.common.geo.GeoPoint getValue()

‎modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/50_script_doc_values.yml

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,56 @@ setup:
132132
- match: { hits.hits.0.fields.field.0.lat: 41.1199999647215 }
133133
- match: { hits.hits.0.fields.field.0.lon: -71.34000004269183 }
134134

135+
- do:
136+
search:
137+
rest_total_hits_as_int: true
138+
body:
139+
script_fields:
140+
centroid:
141+
script:
142+
source: "doc['geo_point'].getCentroid()"
143+
- match: { hits.hits.0.fields.centroid.0.lat: 41.1199999647215 }
144+
- match: { hits.hits.0.fields.centroid.0.lon: -71.34000004269183 }
145+
146+
- do:
147+
search:
148+
rest_total_hits_as_int: true
149+
body:
150+
script_fields:
151+
bbox:
152+
script:
153+
source: "doc['geo_point'].getBoundingBox()"
154+
- match: { hits.hits.0.fields.bbox.0.top_left.lat: 41.1199999647215 }
155+
- match: { hits.hits.0.fields.bbox.0.top_left.lon: -71.34000004269183 }
156+
- match: { hits.hits.0.fields.bbox.0.bottom_right.lat: 41.1199999647215 }
157+
- match: { hits.hits.0.fields.bbox.0.bottom_right.lon: -71.34000004269183 }
158+
159+
- do:
160+
search:
161+
rest_total_hits_as_int: true
162+
body:
163+
script_fields:
164+
topLeft:
165+
script:
166+
source: "doc['geo_point'].getBoundingBox().topLeft()"
167+
bottomRight:
168+
script:
169+
source: "doc['geo_point'].getBoundingBox().bottomRight()"
170+
- match: { hits.hits.0.fields.topLeft.0.lat: 41.1199999647215 }
171+
- match: { hits.hits.0.fields.topLeft.0.lon: -71.34000004269183 }
172+
- match: { hits.hits.0.fields.bottomRight.0.lat: 41.1199999647215 }
173+
- match: { hits.hits.0.fields.bottomRight.0.lon: -71.34000004269183 }
174+
175+
- do:
176+
search:
177+
rest_total_hits_as_int: true
178+
body:
179+
script_fields:
180+
type:
181+
script:
182+
source: "doc['geo_point'].getDimensionalType()"
183+
- match: { hits.hits.0.fields.type.0: 0 }
184+
135185
---
136186
"ip":
137187
- do:

‎server/src/internalClusterTest/java/org/elasticsearch/search/geo/GeoPointScriptDocValuesIT.java

Lines changed: 230 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,230 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
/*
9+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
10+
* or more contributor license agreements. Licensed under the Elastic License
11+
* 2.0 and the Server Side Public License, v 1; you may not use this file except
12+
* in compliance with, at your election, the Elastic License 2.0 or the Server
13+
* Side Public License, v 1.
14+
*/
15+
16+
package org.elasticsearch.search.geo;
17+
18+
import org.apache.lucene.geo.GeoEncodingUtils;
19+
import org.elasticsearch.action.search.SearchResponse;
20+
import org.elasticsearch.common.document.DocumentField;
21+
import org.elasticsearch.common.geo.GeoBoundingBox;
22+
import org.elasticsearch.common.xcontent.XContentBuilder;
23+
import org.elasticsearch.common.xcontent.XContentFactory;
24+
import org.elasticsearch.geo.GeometryTestUtils;
25+
import org.elasticsearch.index.fielddata.ScriptDocValues;
26+
import org.elasticsearch.plugins.Plugin;
27+
import org.elasticsearch.script.MockScriptPlugin;
28+
import org.elasticsearch.script.Script;
29+
import org.elasticsearch.script.ScriptType;
30+
import org.elasticsearch.test.ESSingleNodeTestCase;
31+
import org.hamcrest.Matchers;
32+
import org.junit.Before;
33+
34+
import java.io.IOException;
35+
import java.util.Arrays;
36+
import java.util.Collection;
37+
import java.util.Collections;
38+
import java.util.HashMap;
39+
import java.util.Map;
40+
import java.util.function.Function;
41+
42+
import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
43+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
44+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchResponse;
45+
import static org.hamcrest.Matchers.equalTo;
46+
47+
public class GeoPointScriptDocValuesIT extends ESSingleNodeTestCase {
48+
49+
@Override
50+
protected Collection<Class<? extends Plugin>> getPlugins() {
51+
return Arrays.asList(CustomScriptPlugin.class);
52+
}
53+
54+
public static class CustomScriptPlugin extends MockScriptPlugin {
55+
56+
@Override
57+
protected Map<String, Function<Map<String, Object>, Object>> pluginScripts() {
58+
Map<String, Function<Map<String, Object>, Object>> scripts = new HashMap<>();
59+
60+
scripts.put("lat", this::scriptLat);
61+
scripts.put("lon", this::scriptLon);
62+
scripts.put("height", this::scriptHeight);
63+
scripts.put("width", this::scriptWidth);
64+
return scripts;
65+
}
66+
67+
private double scriptHeight(Map<String, Object> vars) {
68+
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
69+
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
70+
if (geometry.size() == 0) {
71+
return Double.NaN;
72+
} else {
73+
GeoBoundingBox boundingBox = geometry.getBoundingBox();
74+
return boundingBox.topLeft().lat() - boundingBox.bottomRight().lat();
75+
}
76+
}
77+
78+
private double scriptWidth(Map<String, Object> vars) {
79+
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
80+
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
81+
if (geometry.size() == 0) {
82+
return Double.NaN;
83+
} else {
84+
GeoBoundingBox boundingBox = geometry.getBoundingBox();
85+
return boundingBox.bottomRight().lon() - boundingBox.topLeft().lon();
86+
}
87+
}
88+
89+
private double scriptLat(Map<String, Object> vars) {
90+
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
91+
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
92+
return geometry.size() == 0 ? Double.NaN : geometry.getCentroid().lat();
93+
}
94+
95+
private double scriptLon(Map<String, Object> vars) {
96+
Map<?, ?> doc = (Map<?, ?>) vars.get("doc");
97+
ScriptDocValues.Geometry<?> geometry = assertGeometry(doc);
98+
return geometry.size() == 0 ? Double.NaN : geometry.getCentroid().lon();
99+
}
100+
101+
private ScriptDocValues.Geometry<?> assertGeometry(Map<?, ?> doc) {
102+
ScriptDocValues.Geometry<?> geometry = (ScriptDocValues.Geometry<?>) doc.get("location");
103+
if (geometry.size() == 0) {
104+
assertThat(geometry.getBoundingBox(), Matchers.nullValue());
105+
assertThat(geometry.getCentroid(), Matchers.nullValue());
106+
assertThat(geometry.getDimensionalType(), equalTo(-1));
107+
} else {
108+
assertThat(geometry.getBoundingBox(), Matchers.notNullValue());
109+
assertThat(geometry.getCentroid(), Matchers.notNullValue());
110+
assertThat(geometry.getDimensionalType(), equalTo(0));
111+
}
112+
return geometry;
113+
}
114+
}
115+
116+
@Override
117+
protected boolean forbidPrivateIndexSettings() {
118+
return false;
119+
}
120+
121+
@Before
122+
public void setupTestIndex() throws IOException {
123+
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder().startObject().startObject("_doc")
124+
.startObject("properties").startObject("location").field("type", "geo_point");
125+
xContentBuilder.endObject().endObject().endObject().endObject();
126+
assertAcked(client().admin().indices().prepareCreate("test").addMapping("_doc", xContentBuilder));
127+
ensureGreen();
128+
}
129+
130+
public void testRandomPoint() throws Exception {
131+
final double lat = GeometryTestUtils.randomLat();
132+
final double lon = GeometryTestUtils.randomLon();
133+
client().prepareIndex("test", "_doc").setId("1")
134+
.setSource(jsonBuilder().startObject()
135+
.field("name", "TestPosition")
136+
.field("location", new double[]{lon, lat})
137+
.endObject())
138+
.get();
139+
140+
client().admin().indices().prepareRefresh("test").get();
141+
142+
SearchResponse searchResponse = client().prepareSearch().addStoredField("_source")
143+
.addScriptField("lat", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lat", Collections.emptyMap()))
144+
.addScriptField("lon", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lon", Collections.emptyMap()))
145+
.addScriptField("height", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "height", Collections.emptyMap()))
146+
.addScriptField("width", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "width", Collections.emptyMap()))
147+
.get();
148+
assertSearchResponse(searchResponse);
149+
150+
final double qLat = GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(lat));
151+
final double qLon = GeoEncodingUtils.decodeLongitude(GeoEncodingUtils.encodeLongitude(lon));
152+
153+
Map<String, DocumentField> fields = searchResponse.getHits().getHits()[0].getFields();
154+
assertThat(fields.get("lat").getValue(), equalTo(qLat));
155+
assertThat(fields.get("lon").getValue(), equalTo(qLon));
156+
assertThat(fields.get("height").getValue(), equalTo(0d));
157+
assertThat(fields.get("width").getValue(), equalTo(0d));
158+
}
159+
160+
public void testRandomMultiPoint() throws Exception {
161+
final int size = randomIntBetween(2, 20);
162+
final double[] lats = new double[size];
163+
final double[] lons = new double[size];
164+
for (int i = 0; i < size; i++) {
165+
lats[i] = GeometryTestUtils.randomLat();
166+
lons[i] = GeometryTestUtils.randomLon();
167+
}
168+
169+
final double[][] values = new double[size][];
170+
for (int i = 0; i < size; i++) {
171+
values[i] = new double[]{lons[i], lats[i]};
172+
}
173+
174+
XContentBuilder builder = jsonBuilder().startObject()
175+
.field("name", "TestPosition")
176+
.field("location", values).endObject();
177+
client().prepareIndex("test", "_doc").setId("1").setSource(builder).get();
178+
179+
client().admin().indices().prepareRefresh("test").get();
180+
181+
SearchResponse searchResponse = client().prepareSearch().addStoredField("_source")
182+
.addScriptField("lat", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lat", Collections.emptyMap()))
183+
.addScriptField("lon", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lon", Collections.emptyMap()))
184+
.addScriptField("height", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "height", Collections.emptyMap()))
185+
.addScriptField("width", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "width", Collections.emptyMap()))
186+
.get();
187+
assertSearchResponse(searchResponse);
188+
189+
for (int i = 0; i < size; i++) {
190+
lats[i] = GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(lats[i]));
191+
lons[i] = GeoEncodingUtils.decodeLongitude(GeoEncodingUtils.encodeLongitude(lons[i]));
192+
}
193+
194+
final double centroidLon = Arrays.stream(lons).sum() / size;
195+
final double centroidLat = Arrays.stream(lats).sum() / size;
196+
final double width = Arrays.stream(lons).max().getAsDouble() - Arrays.stream(lons).min().getAsDouble();
197+
final double height = Arrays.stream(lats).max().getAsDouble() - Arrays.stream(lats).min().getAsDouble();
198+
199+
Map<String, DocumentField> fields = searchResponse.getHits().getHits()[0].getFields();
200+
assertThat(fields.get("lat").getValue(), equalTo(centroidLat));
201+
assertThat(fields.get("lon").getValue(), equalTo(centroidLon));
202+
assertThat(fields.get("height").getValue(), equalTo(height));
203+
assertThat(fields.get("width").getValue(), equalTo(width));
204+
}
205+
206+
public void testNullPoint() throws Exception {
207+
client().prepareIndex("test", "_doc").setId("1")
208+
.setSource(jsonBuilder().startObject()
209+
.field("name", "TestPosition")
210+
.nullField("location")
211+
.endObject())
212+
.get();
213+
214+
client().admin().indices().prepareRefresh("test").get();
215+
216+
SearchResponse searchResponse = client().prepareSearch().addStoredField("_source")
217+
.addScriptField("lat", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lat", Collections.emptyMap()))
218+
.addScriptField("lon", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "lon", Collections.emptyMap()))
219+
.addScriptField("height", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "height", Collections.emptyMap()))
220+
.addScriptField("width", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "width", Collections.emptyMap()))
221+
.get();
222+
assertSearchResponse(searchResponse);
223+
224+
Map<String, DocumentField> fields = searchResponse.getHits().getHits()[0].getFields();
225+
assertThat(fields.get("lat").getValue(), equalTo(Double.NaN));
226+
assertThat(fields.get("lon").getValue(), equalTo(Double.NaN));
227+
assertThat(fields.get("height").getValue(), equalTo(Double.NaN));
228+
assertThat(fields.get("width").getValue(), equalTo(Double.NaN));
229+
}
230+
}

‎server/src/main/java/org/elasticsearch/common/geo/GeoBoundingBox.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import org.elasticsearch.common.io.stream.StreamInput;
1313
import org.elasticsearch.common.io.stream.StreamOutput;
1414
import org.elasticsearch.common.io.stream.Writeable;
15-
import org.elasticsearch.common.xcontent.ToXContentObject;
15+
import org.elasticsearch.common.xcontent.ToXContentFragment;
1616
import org.elasticsearch.common.xcontent.XContentBuilder;
1717
import org.elasticsearch.common.xcontent.XContentParser;
1818
import org.elasticsearch.geometry.Geometry;
@@ -29,7 +29,7 @@
2929
* A class representing a Geo-Bounding-Box for use by Geo queries and aggregations
3030
* that deal with extents/rectangles representing rectangular areas of interest.
3131
*/
32-
public class GeoBoundingBox implements ToXContentObject, Writeable {
32+
public class GeoBoundingBox implements ToXContentFragment, Writeable {
3333
private static final WellKnownText WKT_PARSER = new WellKnownText(true, new StandardValidator(true));
3434
static final ParseField TOP_RIGHT_FIELD = new ParseField("top_right");
3535
static final ParseField BOTTOM_LEFT_FIELD = new ParseField("bottom_left");
@@ -88,7 +88,7 @@ public double right() {
8888

8989
@Override
9090
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
91-
builder.startObject(BOUNDS_FIELD.getPreferredName());
91+
builder.startObject();
9292
toXContentFragment(builder, true);
9393
builder.endObject();
9494
return builder;

0 commit comments

Comments
 (0)