Skip to content

Commit

Permalink
Support for geo_distance queries on geo_shape fields (#2516)
Browse files Browse the repository at this point in the history
Adds support for using the geo_distance query on geo_shape field types by
lifting the geo_point restriction in the QueryBuilder. Distance query
integration tests are abstracted to test both geo_point and geo_shape field
types.

Signed-off-by: Nicholas Walter Knize <[email protected]>
  • Loading branch information
nknize authored Mar 21, 2022
1 parent 42aa7ab commit 4bb8add
Show file tree
Hide file tree
Showing 5 changed files with 194 additions and 38 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -44,18 +44,19 @@
import org.opensearch.common.xcontent.XContentFactory;
import org.opensearch.geometry.utils.Geohash;
import org.opensearch.index.fielddata.ScriptDocValues;
import org.opensearch.index.query.IdsQueryBuilder;
import org.opensearch.index.query.QueryBuilders;
import org.opensearch.plugins.Plugin;
import org.opensearch.script.MockScriptPlugin;
import org.opensearch.script.Script;
import org.opensearch.script.ScriptType;
import org.opensearch.search.SearchHit;
import org.opensearch.search.aggregations.AggregationBuilders;
import org.opensearch.search.aggregations.Aggregations;
import org.opensearch.search.aggregations.bucket.range.InternalGeoDistance;
import org.opensearch.search.aggregations.bucket.range.Range;
import org.opensearch.test.OpenSearchIntegTestCase;
import org.opensearch.test.VersionUtils;
import org.junit.Before;

import java.io.IOException;
import java.util.Collection;
Expand All @@ -65,11 +66,14 @@
import java.util.Map;
import java.util.function.Function;

import static org.hamcrest.Matchers.anyOf;
import static org.hamcrest.Matchers.equalTo;
import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
import static org.hamcrest.Matchers.closeTo;

public class GeoDistanceIT extends OpenSearchIntegTestCase {
/** base class for testing geo_distance queries on geo_ field types */
abstract class AbstractGeoDistanceIT extends OpenSearchIntegTestCase {

private static final double src_lat = 32.798;
private static final double src_lon = -117.151;
Expand Down Expand Up @@ -118,30 +122,89 @@ protected boolean forbidPrivateIndexSettings() {
return false;
}

@Before
public void setupTestIndex() throws IOException {
protected void indexSetup() throws IOException {
Version version = VersionUtils.randomIndexCompatibleVersion(random());
Settings settings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, version).build();
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder()
.startObject()
.startObject("properties")
.startObject("location")
.field("type", "geo_point");
xContentBuilder.endObject().endObject().endObject();
XContentBuilder xContentBuilder = getMapping();
assertAcked(prepareCreate("test").setSettings(settings).setMapping(xContentBuilder));
ensureGreen();

client().prepareIndex("test")
.setId("1")
.setSource(jsonBuilder().startObject().field("name", "New York").field("location", "POINT(-74.0059731 40.7143528)").endObject())
.get();

// to NY: 5.286 km
client().prepareIndex("test")
.setId("2")
.setSource(
jsonBuilder().startObject().field("name", "Times Square").field("location", "POINT(-73.9844722 40.759011)").endObject()
)
.get();

// to NY: 0.4621 km
client().prepareIndex("test")
.setId("3")
.setSource(jsonBuilder().startObject().field("name", "Tribeca").field("location", "POINT(-74.007819 40.718266)").endObject())
.get();

// to NY: 1.055 km
client().prepareIndex("test")
.setId("4")
.setSource(
jsonBuilder().startObject().field("name", "Wall Street").field("location", "POINT(-74.0088305 40.7051157)").endObject()
)
.get();

// to NY: 1.258 km
client().prepareIndex("test")
.setId("5")
.setSource(jsonBuilder().startObject().field("name", "Soho").field("location", "POINT(-74 40.7247222)").endObject())
.get();

// to NY: 2.029 km
client().prepareIndex("test")
.setId("6")
.setSource(
jsonBuilder().startObject().field("name", "Greenwich Village").field("location", "POINT(-73.9962255 40.731033)").endObject()
)
.get();

// to NY: 8.572 km
client().prepareIndex("test")
.setId("7")
.setSource(jsonBuilder().startObject().field("name", "Brooklyn").field("location", "POINT(-73.95 40.65)").endObject())
.get();

client().admin().indices().prepareRefresh().get();
}

public abstract XContentBuilder addGeoMapping(XContentBuilder parentMapping) throws IOException;

public XContentBuilder getMapping() throws IOException {
XContentBuilder mapping = XContentFactory.jsonBuilder().startObject().startObject("properties");
mapping = addGeoMapping(mapping);
return mapping.endObject().endObject();
}

public void testSimpleDistanceQuery() {
SearchResponse searchResponse = client().prepareSearch() // from NY
.setQuery(QueryBuilders.geoDistanceQuery("location").point(40.5, -73.9).distance(25, DistanceUnit.KILOMETERS))
.get();
assertThat(searchResponse.getHits().getTotalHits().value, equalTo(2L));
assertThat(searchResponse.getHits().getHits().length, equalTo(2));
for (SearchHit hit : searchResponse.getHits()) {
assertThat(hit.getId(), anyOf(equalTo("7"), equalTo("4")));
}
}

public void testDistanceScript() throws Exception {
client().prepareIndex("test")
.setId("1")
.setId("8")
.setSource(
jsonBuilder().startObject()
.field("name", "TestPosition")
.startObject("location")
.field("lat", src_lat)
.field("lon", src_lon)
.endObject()
.field("location", "POINT(" + src_lon + " " + src_lat + ")")
.endObject()
)
.get();
Expand All @@ -150,6 +213,7 @@ public void testDistanceScript() throws Exception {

// Test doc['location'].arcDistance(lat, lon)
SearchResponse searchResponse1 = client().prepareSearch()
.setQuery(new IdsQueryBuilder().addIds("8"))
.addStoredField("_source")
.addScriptField("distance", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "arcDistance", Collections.emptyMap()))
.get();
Expand All @@ -158,6 +222,7 @@ public void testDistanceScript() throws Exception {

// Test doc['location'].planeDistance(lat, lon)
SearchResponse searchResponse2 = client().prepareSearch()
.setQuery(new IdsQueryBuilder().addIds("8"))
.addStoredField("_source")
.addScriptField("distance", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "planeDistance", Collections.emptyMap()))
.get();
Expand All @@ -166,6 +231,7 @@ public void testDistanceScript() throws Exception {

// Test doc['location'].geohashDistance(lat, lon)
SearchResponse searchResponse4 = client().prepareSearch()
.setQuery(new IdsQueryBuilder().addIds("8"))
.addStoredField("_source")
.addScriptField("distance", new Script(ScriptType.INLINE, CustomScriptPlugin.NAME, "geohashDistance", Collections.emptyMap()))
.get();
Expand All @@ -180,6 +246,7 @@ public void testDistanceScript() throws Exception {

// Test doc['location'].arcDistance(lat, lon + 360)/1000d
SearchResponse searchResponse5 = client().prepareSearch()
.setQuery(new IdsQueryBuilder().addIds("8"))
.addStoredField("_source")
.addScriptField(
"distance",
Expand All @@ -191,6 +258,7 @@ public void testDistanceScript() throws Exception {

// Test doc['location'].arcDistance(lat + 360, lon)/1000d
SearchResponse searchResponse6 = client().prepareSearch()
.setQuery(new IdsQueryBuilder().addIds("8"))
.addStoredField("_source")
.addScriptField(
"distance",
Expand All @@ -207,10 +275,7 @@ public void testGeoDistanceAggregation() throws IOException {
.setSource(
jsonBuilder().startObject()
.field("name", "TestPosition")
.startObject("location")
.field("lat", src_lat)
.field("lon", src_lon)
.endObject()
.field("location", "POINT(" + src_lon + " " + src_lat + ")")
.endObject()
)
.get();
Expand All @@ -225,7 +290,7 @@ public void testGeoDistanceAggregation() throws IOException {
AggregationBuilders.geoDistance(name, new GeoPoint(tgt_lat, tgt_lon))
.field("location")
.unit(DistanceUnit.MILES)
.addRange(0, 25000)
.addRange(0, 25000) // limits the distance (expected to omit one point outside this range)
);

search.setSize(0); // no hits please
Expand All @@ -239,6 +304,6 @@ public void testGeoDistanceAggregation() throws IOException {
List<? extends Range.Bucket> buckets = ((Range) geoDistance).getBuckets();
assertNotNull("Buckets should not be null", buckets);
assertEquals("Unexpected number of buckets", 1, buckets.size());
assertEquals("Unexpected doc count for geo distance", 1, buckets.get(0).getDocCount());
assertEquals("Unexpected doc count for geo distance", 7, buckets.get(0).getDocCount());
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package org.opensearch.search.geo;

import org.junit.Before;
import org.opensearch.common.xcontent.XContentBuilder;

import java.io.IOException;

/** Tests geo_distance queries on geo_point field types */
public class GeoDistanceQueryGeoPointsIT extends AbstractGeoDistanceIT {

@Before
public void setupTestIndex() throws IOException {
indexSetup();
}

@Override
public XContentBuilder addGeoMapping(XContentBuilder parentMapping) throws IOException {
return parentMapping.startObject("location").field("type", "geo_point").endObject();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package org.opensearch.search.geo;

import org.junit.Before;
import org.opensearch.common.xcontent.XContentBuilder;

import java.io.IOException;

/** Tests geo_distance queries on geo_shape field types */
public class GeoDistanceQueryGeoShapesIT extends AbstractGeoDistanceIT {

@Before
public void setupTestIndex() throws IOException {
indexSetup();
}

@Override
public XContentBuilder addGeoMapping(XContentBuilder parentMapping) throws IOException {
parentMapping = parentMapping.startObject("location").field("type", "geo_shape");
if (randomBoolean()) {
parentMapping.field("strategy", "recursive");
}
return parentMapping.endObject();
}

@AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/2515")
@Override
public void testDistanceScript() {
// no-op; todo script support for distance calculation on shapes cannot be added until GeoShapeDocValues is added
}

@AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/2515")
@Override
public void testGeoDistanceAggregation() {
// no-op; todo aggregation support for distance calculation on shapes cannot be added until GeoShapeDocValues is added
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,6 @@

package org.opensearch.index.query;

import org.apache.lucene.document.LatLonDocValuesField;
import org.apache.lucene.document.LatLonPoint;
import org.apache.lucene.search.IndexOrDocValuesQuery;
import org.apache.lucene.search.MatchNoDocsQuery;
import org.apache.lucene.search.Query;
import org.opensearch.common.ParseField;
Expand All @@ -43,12 +40,17 @@
import org.opensearch.common.geo.GeoDistance;
import org.opensearch.common.geo.GeoPoint;
import org.opensearch.common.geo.GeoUtils;
import org.opensearch.common.geo.ShapeRelation;
import org.opensearch.common.geo.SpatialStrategy;
import org.opensearch.common.io.stream.StreamInput;
import org.opensearch.common.io.stream.StreamOutput;
import org.opensearch.common.unit.DistanceUnit;
import org.opensearch.common.xcontent.XContentBuilder;
import org.opensearch.common.xcontent.XContentParser;
import org.opensearch.index.mapper.GeoPointFieldMapper.GeoPointFieldType;
import org.opensearch.geometry.Circle;
import org.opensearch.index.mapper.GeoPointFieldMapper;
import org.opensearch.index.mapper.GeoShapeFieldMapper;
import org.opensearch.index.mapper.GeoShapeQueryable;
import org.opensearch.index.mapper.MappedFieldType;

import java.io.IOException;
Expand Down Expand Up @@ -245,12 +247,25 @@ protected Query doToQuery(QueryShardContext shardContext) throws IOException {
if (ignoreUnmapped) {
return new MatchNoDocsQuery();
} else {
throw new QueryShardException(shardContext, "failed to find geo_point field [" + fieldName + "]");
throw new QueryShardException(shardContext, "failed to find geo field [" + fieldName + "]");
}
}

if (!(fieldType instanceof GeoPointFieldType)) {
throw new QueryShardException(shardContext, "field [" + fieldName + "] is not a geo_point field");
if (fieldType instanceof GeoShapeQueryable == false) {
throw new QueryShardException(
shardContext,
"type ["
+ fieldType
+ "] for field ["
+ fieldName
+ "] is not supported for ["
+ NAME
+ "] queries. Must be one of ["
+ GeoPointFieldMapper.CONTENT_TYPE
+ "] or ["
+ GeoShapeFieldMapper.CONTENT_TYPE
+ "]"
);
}

QueryValidationException exception = checkLatLon();
Expand All @@ -262,12 +277,9 @@ protected Query doToQuery(QueryShardContext shardContext) throws IOException {
GeoUtils.normalizePoint(center, true, true);
}

Query query = LatLonPoint.newDistanceQuery(fieldType.name(), center.lat(), center.lon(), this.distance);
if (fieldType.hasDocValues()) {
Query dvQuery = LatLonDocValuesField.newSlowDistanceQuery(fieldType.name(), center.lat(), center.lon(), this.distance);
query = new IndexOrDocValuesQuery(query, dvQuery);
}
return query;
final GeoShapeQueryable geoShapeQueryable = (GeoShapeQueryable) fieldType;
final Circle circle = new Circle(center.lon(), center.lat(), this.distance);
return geoShapeQueryable.geoShapeQuery(circle, fieldType.name(), SpatialStrategy.RECURSIVE, ShapeRelation.INTERSECTS, shardContext);
}

@Override
Expand Down
Loading

0 comments on commit 4bb8add

Please sign in to comment.