Support for geo_distance queries on geo_shape fields (#2516)

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 <nknize@apache.org>
This commit is contained in:
Nick Knize 2022-03-21 18:47:49 -05:00 committed by GitHub
parent 42aa7ab68a
commit 4bb8add982
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 194 additions and 38 deletions

View File

@ -44,18 +44,19 @@ import org.opensearch.common.xcontent.XContentBuilder;
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;
@ -65,11 +66,14 @@ import java.util.List;
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;
@ -118,30 +122,89 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
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();
@ -150,6 +213,7 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
// 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();
@ -158,6 +222,7 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
// 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();
@ -166,6 +231,7 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
// 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();
@ -180,6 +246,7 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
// Test doc['location'].arcDistance(lat, lon + 360)/1000d
SearchResponse searchResponse5 = client().prepareSearch()
.setQuery(new IdsQueryBuilder().addIds("8"))
.addStoredField("_source")
.addScriptField(
"distance",
@ -191,6 +258,7 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
// Test doc['location'].arcDistance(lat + 360, lon)/1000d
SearchResponse searchResponse6 = client().prepareSearch()
.setQuery(new IdsQueryBuilder().addIds("8"))
.addStoredField("_source")
.addScriptField(
"distance",
@ -207,10 +275,7 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
.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();
@ -225,7 +290,7 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
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
@ -239,6 +304,6 @@ public class GeoDistanceIT extends OpenSearchIntegTestCase {
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());
}
}

View File

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

View File

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

View File

@ -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;
@ -43,12 +40,17 @@ import org.opensearch.common.Strings;
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;
@ -245,12 +247,25 @@ public class GeoDistanceQueryBuilder extends AbstractQueryBuilder<GeoDistanceQue
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();
@ -262,12 +277,9 @@ public class GeoDistanceQueryBuilder extends AbstractQueryBuilder<GeoDistanceQue
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

View File

@ -41,6 +41,9 @@ import org.opensearch.common.ParsingException;
import org.opensearch.common.geo.GeoDistance;
import org.opensearch.common.geo.GeoPoint;
import org.opensearch.common.unit.DistanceUnit;
import org.opensearch.index.mapper.GeoPointFieldMapper;
import org.opensearch.index.mapper.GeoShapeFieldMapper;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.test.AbstractQueryTestCase;
import org.opensearch.test.geo.RandomShapeGenerator;
import org.locationtech.spatial4j.shape.Point;
@ -55,7 +58,7 @@ public class GeoDistanceQueryBuilderTests extends AbstractQueryTestCase<GeoDista
@Override
protected GeoDistanceQueryBuilder doCreateTestQueryBuilder() {
String fieldName = randomFrom(GEO_POINT_FIELD_NAME, GEO_POINT_ALIAS_FIELD_NAME);
String fieldName = randomFrom(GEO_POINT_FIELD_NAME, GEO_POINT_ALIAS_FIELD_NAME, GEO_SHAPE_FIELD_NAME);
GeoDistanceQueryBuilder qb = new GeoDistanceQueryBuilder(fieldName);
String distance = "" + randomDouble();
if (randomBoolean()) {
@ -141,8 +144,10 @@ public class GeoDistanceQueryBuilderTests extends AbstractQueryTestCase<GeoDista
@Override
protected void doAssertLuceneQuery(GeoDistanceQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
// TODO: remove the if statement once we always use LatLonPoint
if (query instanceof IndexOrDocValuesQuery) {
final MappedFieldType fieldType = context.getFieldType(queryBuilder.fieldName());
if (fieldType == null) {
assertTrue("Found no indexed geo query.", query instanceof MatchNoDocsQuery);
} else if (fieldType instanceof GeoPointFieldMapper.GeoPointFieldType) {
Query indexQuery = ((IndexOrDocValuesQuery) query).getIndexQuery();
String expectedFieldName = expectedFieldName(queryBuilder.fieldName());
@ -165,6 +170,8 @@ public class GeoDistanceQueryBuilderTests extends AbstractQueryTestCase<GeoDista
),
dvQuery
);
} else {
assertEquals(GeoShapeFieldMapper.GeoShapeFieldType.class, fieldType.getClass());
}
}
@ -382,7 +389,7 @@ public class GeoDistanceQueryBuilderTests extends AbstractQueryTestCase<GeoDista
final GeoDistanceQueryBuilder failingQueryBuilder = new GeoDistanceQueryBuilder("unmapped").point(0.0, 0.0).distance("20m");
failingQueryBuilder.ignoreUnmapped(false);
QueryShardException e = expectThrows(QueryShardException.class, () -> failingQueryBuilder.toQuery(shardContext));
assertThat(e.getMessage(), containsString("failed to find geo_point field [unmapped]"));
assertThat(e.getMessage(), containsString("failed to find geo field [unmapped]"));
}
public void testParseFailsWithMultipleFields() throws IOException {