-
Notifications
You must be signed in to change notification settings - Fork 24.4k
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
Changes from 3 commits
d13edc0
7eb560e
4febdc3
f488c94
8be69d0
31d6b9b
b019a47
b994a04
35154e7
1de954e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
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 |
---|---|---|
|
@@ -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; | ||
|
@@ -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) { | ||
|
@@ -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); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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. There was a problem hiding this comment. Choose a reason for hiding this commentThe 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); | ||
} | ||
|
@@ -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(); | ||
} | ||
} | ||
} | ||
|
@@ -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> { | ||
|
There was a problem hiding this comment.
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