diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/Component2DVisitor.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/Component2DVisitor.java new file mode 100644 index 0000000000000..fe9b4e0dff6e3 --- /dev/null +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/Component2DVisitor.java @@ -0,0 +1,388 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.spatial.index.fielddata; + +import org.apache.lucene.document.ShapeField; +import org.apache.lucene.geo.Component2D; +import org.apache.lucene.index.PointValues; + +/** + * A {@link TriangleTreeReader.Visitor} implementation for {@link Component2D} geometries. + * It can solve spatial relationships against a serialize triangle tree. + */ +public abstract class Component2DVisitor implements TriangleTreeReader.Visitor { + + protected final Component2D component2D; + private final CoordinateEncoder encoder; + + private Component2DVisitor(Component2D component2D, CoordinateEncoder encoder) { + this.component2D = component2D; + this.encoder = encoder; + } + + /** If the relationship has been honour. */ + public abstract boolean matches(); + + /** Reset the visitor to the initial state. */ + public abstract void reset(); + + @Override + public void visitPoint(int x, int y) { + doVisitPoint(encoder.decodeX(x), encoder.decodeY(y)); + } + + abstract void doVisitPoint(double x, double y); + + @Override + public void visitLine(int aX, int aY, int bX, int bY, byte metadata) { + doVisitLine(encoder.decodeX(aX), encoder.decodeY(aY), encoder.decodeX(bX), encoder.decodeY(bY), metadata); + } + + abstract void doVisitLine(double aX, double aY, double bX, double bY, byte metadata); + + @Override + public void visitTriangle(int aX, int aY, int bX, int bY, int cX, int cY, byte metadata) { + doVisitTriangle( + encoder.decodeX(aX), + encoder.decodeY(aY), + encoder.decodeX(bX), + encoder.decodeY(bY), + encoder.decodeX(cX), + encoder.decodeY(cY), + metadata + ); + } + + abstract void doVisitTriangle(double aX, double aY, double bX, double bY, double cX, double cY, byte metadata); + + @Override + public boolean pushX(int minX) { + return component2D.getMaxX() >= encoder.decodeX(minX); + } + + @Override + public boolean pushY(int minY) { + return component2D.getMaxY() >= encoder.decodeY(minY); + } + + @Override + public boolean push(int maxX, int maxY) { + return component2D.getMinX() <= encoder.decodeX(maxX) && + component2D.getMinY() <= encoder.decodeY(maxY); + + } + + @Override + public boolean push(int minX, int minY, int maxX, int maxY) { + final PointValues.Relation relation = component2D.relate( + encoder.decodeX(minX), + encoder.decodeX(maxX), + encoder.decodeY(minY), + encoder.decodeY(maxY) + ); + return doPush(relation); + } + + /** Relation between the query shape and the doc value bounding box. Depending on the query relationship, + * decide if we should traverse the tree. + * + * @return if true, the visitor keeps traversing the tree, else it stops. + * */ + abstract boolean doPush(PointValues.Relation relation); + + /** + * Creates a visitor from the provided Component2D and spatial relationship. Visitors are re-usable by + * calling the {@link #reset()} method. + */ + public static Component2DVisitor getVisitor( + Component2D component2D, + ShapeField.QueryRelation relation, + CoordinateEncoder encoder + ) { + switch (relation) { + case CONTAINS: + return new ContainsVisitor(component2D, encoder); + case INTERSECTS: + return new IntersectsVisitor(component2D, encoder); + case DISJOINT: + return new DisjointVisitor(component2D, encoder); + case WITHIN: + return new WithinVisitor(component2D, encoder); + default: + throw new IllegalArgumentException("Invalid query relation:[" + relation + "]"); + } + } + + /** + * Intersects visitor stops as soon as there is one triangle intersecting the component + */ + private static class IntersectsVisitor extends Component2DVisitor { + + boolean intersects; + + private IntersectsVisitor(Component2D component2D, CoordinateEncoder encoder) { + super(component2D, encoder); + } + + @Override + public boolean matches() { + return intersects; + } + + @Override + public void reset() { + // Start assuming that shapes are disjoint. As soon an intersecting component is found, + // stop traversing the tree. + intersects = false; + } + + @Override + void doVisitPoint(double x, double y) { + intersects = component2D.contains(x, y); + } + + @Override + void doVisitLine(double aX, double aY, double bX, double bY, byte metadata) { + intersects = component2D.intersectsLine(aX, aY, bX, bY); + } + + @Override + void doVisitTriangle(double aX, double aY, double bX, double bY, double cX, double cY, byte metadata) { + intersects = component2D.intersectsTriangle(aX, aY, bX, bY, cX, cY); + } + + @Override + public boolean push() { + // as far as shapes don't intersect, keep traversing the tree + return intersects == false; + } + + @Override + boolean doPush(PointValues.Relation relation) { + if (relation == PointValues.Relation.CELL_OUTSIDE_QUERY) { + // shapes are disjoint, stop traversing the tree. + return false; + } else if (relation == PointValues.Relation.CELL_INSIDE_QUERY) { + // shapes intersects, stop traversing the tree. + intersects = true; + return false; + } else { + // traverse the tree. + return true; + } + } + } + + /** + * Disjoint visitor stops as soon as there is one triangle intersecting the component + */ + private static class DisjointVisitor extends Component2DVisitor { + + boolean disjoint; + + private DisjointVisitor(Component2D component2D, CoordinateEncoder encoder) { + super(component2D, encoder); + disjoint = true; + } + + @Override + public boolean matches() { + return disjoint; + } + + @Override + public void reset() { + // Start assuming that shapes are disjoint. As soon an intersecting component is found, + // stop traversing the tree. + disjoint = true; + } + + @Override + void doVisitPoint(double x, double y) { + disjoint = component2D.contains(x, y) == false; + } + + @Override + void doVisitLine(double aX, double aY, double bX, double bY, byte metadata) { + disjoint = component2D.intersectsLine(aX, aY, bX, bY) == false; + } + + @Override + void doVisitTriangle(double aX, double aY, double bX, double bY, double cX, double cY, byte metadata) { + disjoint = component2D.intersectsTriangle(aX, aY, bX, bY, cX, cY) == false; + } + + @Override + public boolean push() { + // as far as the shapes are disjoint, keep traversing the tree + return disjoint; + } + + @Override + boolean doPush(PointValues.Relation relation) { + if (relation == PointValues.Relation.CELL_OUTSIDE_QUERY) { + // shapes are disjoint, stop traversing the tree. + return false; + } else if (relation == PointValues.Relation.CELL_INSIDE_QUERY) { + // shapes intersects, stop traversing the tree. + disjoint = false; + return false; + } else { + // trasverse the tree + return true; + } + } + } + + /** + * within visitor stops as soon as there is one triangle that is not within the component + */ + private static class WithinVisitor extends Component2DVisitor { + + boolean within; + + private WithinVisitor(Component2D component2D, CoordinateEncoder encoder) { + super(component2D, encoder); + within = true; + } + + @Override + public boolean matches() { + return within; + } + + @Override + public void reset() { + // Start assuming that the doc value is within the query shape. As soon + // as a component is not within the query, stop traversing the tree. + within = true; + } + + @Override + void doVisitPoint(double x, double y) { + within = component2D.contains(x, y); + } + + @Override + void doVisitLine(double aX, double aY, double bX, double bY, byte metadata) { + within = component2D.containsLine(aX, aY, bX, bY); + } + + @Override + void doVisitTriangle(double aX, double aY, double bX, double bY, double cX, double cY, byte metadata) { + within = component2D.containsTriangle(aX, aY, bX, bY, cX, cY); + } + + @Override + public boolean push() { + // as far as the doc value is within the query shape, keep traversing the tree + return within; + } + + @Override + public boolean pushX(int minX) { + // if any part of the tree is skipped, then the doc value is not within the shape, + // stop traversing the tree + within = super.pushX(minX); + return within; + } + + @Override + public boolean pushY(int minY) { + // if any part of the tree is skipped, then the doc value is not within the shape, + // stop traversing the tree + within = super.pushY(minY); + return within; + } + + @Override + public boolean push(int maxX, int maxY) { + // if any part of the tree is skipped, then the doc value is not within the shape, + // stop traversing the tree + within = super.push(maxX, maxY); + return within; + } + + @Override + boolean doPush(PointValues.Relation relation) { + if (relation == PointValues.Relation.CELL_OUTSIDE_QUERY) { + // shapes are disjoint, stop traversing the tree. + within = false; + } + return within; + } + } + + /** + * contains visitor stops as soon as there is one triangle that intersects the component + * with an edge belonging to the original polygon. + */ + private static class ContainsVisitor extends Component2DVisitor { + + Component2D.WithinRelation answer; + + private ContainsVisitor(Component2D component2D, CoordinateEncoder encoder) { + super(component2D, encoder); + answer = Component2D.WithinRelation.DISJOINT; + } + + @Override + public boolean matches() { + return answer == Component2D.WithinRelation.CANDIDATE; + } + + @Override + public void reset() { + // Start assuming that shapes are disjoint. As soon + // as a component has a NOTWITHIN relationship, stop traversing the tree. + answer = Component2D.WithinRelation.DISJOINT; + } + + @Override + void doVisitPoint(double x, double y) { + final Component2D.WithinRelation rel = component2D.withinPoint(x, y); + if (rel != Component2D.WithinRelation.DISJOINT) { + // Only override relationship if different to DISJOINT + answer = rel; + } + } + + @Override + void doVisitLine(double aX, double aY, double bX, double bY, byte metadata) { + final boolean ab = (metadata & 1 << 4) == 1 << 4; + final Component2D.WithinRelation rel = component2D.withinLine(aX, aY, ab, bX, bY); + if (rel != Component2D.WithinRelation.DISJOINT) { + // Only override relationship if different to DISJOINT + answer = rel; + } + } + + @Override + void doVisitTriangle(double aX, double aY, double bX, double bY, double cX, double cY, byte metadata) { + final boolean ab = (metadata & 1 << 4) == 1 << 4; + final boolean bc = (metadata & 1 << 5) == 1 << 5; + final boolean ca = (metadata & 1 << 6) == 1 << 6; + final Component2D.WithinRelation rel = component2D.withinTriangle(aX, aY, ab, bX, bY, bc, cX, cY, ca); + if (rel != Component2D.WithinRelation.DISJOINT) { + // Only override relationship if different to DISJOINT + answer = rel; + } + } + + @Override + public boolean push() { + // If the relationship is NOTWITHIN, stop traversing the tree + return answer != Component2D.WithinRelation.NOTWITHIN; + } + + @Override + boolean doPush(PointValues.Relation relation) { + // Only traverse the tree if the shapes intersects. + return relation == PointValues.Relation.CELL_CROSSES_QUERY; + } + } +} diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeometryDocValueReader.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeometryDocValueReader.java index eb62897323d96..04c0f4e6bad4f 100644 --- a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeometryDocValueReader.java +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeometryDocValueReader.java @@ -95,7 +95,7 @@ protected double getSumCentroidWeight() { /** * Visit the triangle tree with the provided visitor */ - protected void visit(TriangleTreeReader.Visitor visitor) { + public void visit(TriangleTreeReader.Visitor visitor) { Extent extent = getExtent(); int thisMaxX = extent.maxX(); int thisMinX = extent.minX(); @@ -105,4 +105,5 @@ protected void visit(TriangleTreeReader.Visitor visitor) { TriangleTreeReader.visit(input, visitor, thisMaxX, thisMaxY); } } + } diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java index fe44deebda005..1ad06c6a2fcba 100644 --- a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java @@ -27,9 +27,9 @@ import org.elasticsearch.index.mapper.Mapper; import org.elasticsearch.index.mapper.ParseContext; import org.elasticsearch.index.query.SearchExecutionContext; -import org.elasticsearch.index.query.VectorGeoShapeQueryProcessor; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.xpack.spatial.index.fielddata.plain.AbstractLatLonShapeIndexFieldData; +import org.elasticsearch.xpack.spatial.index.query.VectorGeoShapeWithDocValuesQueryProcessor; import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType; import java.util.Arrays; @@ -128,7 +128,7 @@ protected void addDocValuesFields(String name, Geometry shape, List meta) { @@ -147,7 +147,7 @@ public String typeName() { @Override public Query geoShapeQuery(Geometry shape, String fieldName, ShapeRelation relation, SearchExecutionContext context) { - return queryProcessor.geoShapeQuery(shape, fieldName, relation, context); + return queryProcessor.geoShapeQuery(shape, fieldName, relation, context, hasDocValues()); } } diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/LatLonShapeDocValuesQuery.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/LatLonShapeDocValuesQuery.java new file mode 100644 index 0000000000000..d1711fcc40201 --- /dev/null +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/LatLonShapeDocValuesQuery.java @@ -0,0 +1,193 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.spatial.index.query; + +import org.apache.lucene.document.ShapeField; +import org.apache.lucene.geo.Component2D; +import org.apache.lucene.geo.LatLonGeometry; +import org.apache.lucene.geo.Rectangle; +import org.apache.lucene.index.BinaryDocValues; +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.ConstantScoreScorer; +import org.apache.lucene.search.ConstantScoreWeight; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.TwoPhaseIterator; +import org.apache.lucene.search.Weight; +import org.elasticsearch.xpack.spatial.index.fielddata.Component2DVisitor; +import org.elasticsearch.xpack.spatial.index.fielddata.CoordinateEncoder; +import org.elasticsearch.xpack.spatial.index.fielddata.GeometryDocValueReader; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** Lucene geometry query for {@link org.elasticsearch.xpack.spatial.index.mapper.BinaryGeoShapeDocValuesField}. */ +class LatLonShapeDocValuesQuery extends Query { + + private final String field; + private final LatLonGeometry[] geometries; + private final ShapeField.QueryRelation relation; + + LatLonShapeDocValuesQuery(String field, ShapeField.QueryRelation relation, LatLonGeometry... geometries) { + if (field == null) { + throw new IllegalArgumentException("field must not be null"); + } + this.field = field; + this.geometries = geometries; + this.relation = relation; + } + + @Override + public String toString(String field) { + StringBuilder sb = new StringBuilder(); + if (!this.field.equals(field)) { + sb.append(this.field); + sb.append(':'); + sb.append(relation); + sb.append(':'); + } + sb.append(Arrays.toString(geometries)); + return sb.toString(); + } + + @Override + public boolean equals(Object obj) { + if (sameClassAs(obj) == false) { + return false; + } + LatLonShapeDocValuesQuery other = (LatLonShapeDocValuesQuery) obj; + return field.equals(other.field) && relation == other.relation && Arrays.equals(geometries, other.geometries); + } + + @Override + public int hashCode() { + int h = classHash(); + h = 31 * h + field.hashCode(); + h = 31 * h + relation.hashCode(); + h = 31 * h + Arrays.hashCode(geometries); + return h; + } + + @Override + public void visit(QueryVisitor visitor) { + if (visitor.acceptField(field)) { + visitor.visitLeaf(this); + } + } + + @Override + public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) { + if (relation == ShapeField.QueryRelation.CONTAINS) { + return getContainsWeight(scoreMode, boost); + } else { + return getStandardWeight(scoreMode, boost); + } + } + + private ConstantScoreWeight getStandardWeight(ScoreMode scoreMode, float boost) { + return new ConstantScoreWeight(this, boost) { + final Component2D component2D = LatLonGeometry.create(geometries); + + @Override + public Scorer scorer(LeafReaderContext context) throws IOException { + final BinaryDocValues values = context.reader().getBinaryDocValues(field); + if (values == null) { + return null; + } + final GeometryDocValueReader reader = new GeometryDocValueReader(); + final Component2DVisitor visitor = Component2DVisitor.getVisitor(component2D, relation, CoordinateEncoder.GEO); + + final TwoPhaseIterator iterator = new TwoPhaseIterator(values) { + + @Override + public boolean matches() throws IOException { + reader.reset(values.binaryValue()); + visitor.reset(); + reader.visit(visitor); + return visitor.matches(); + } + + @Override + public float matchCost() { + return 1000f; // TODO: what should it be? + } + }; + return new ConstantScoreScorer(this, boost, scoreMode, iterator); + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return DocValues.isCacheable(ctx, field); + } + + }; + } + + private ConstantScoreWeight getContainsWeight(ScoreMode scoreMode, float boost) { + final List components2D = new ArrayList<>(geometries.length); + for (int i = 0; i < geometries.length; i++) { + LatLonGeometry geometry = geometries[i]; + if (geometry instanceof Rectangle) { + Rectangle r = (Rectangle) geometry; + if (r.minLon > r.maxLon) { + components2D.add(LatLonGeometry.create(new Rectangle(r.minLat, r.maxLat, r.minLon, 180))); + components2D.add(LatLonGeometry.create(new Rectangle(r.minLat, r.maxLat, -180, r.maxLon))); + continue; + } + } + components2D.add(LatLonGeometry.create(geometry)); + } + return new ConstantScoreWeight(this, boost) { + + @Override + public Scorer scorer(LeafReaderContext context) throws IOException { + final BinaryDocValues values = context.reader().getBinaryDocValues(field); + if (values == null) { + return null; + } + final GeometryDocValueReader reader = new GeometryDocValueReader(); + final Component2DVisitor[] visitors = new Component2DVisitor[components2D.size()]; + for (int i = 0; i < components2D.size(); i++) { + visitors[i] = Component2DVisitor.getVisitor(components2D.get(i), relation, CoordinateEncoder.GEO); + } + + final TwoPhaseIterator iterator = new TwoPhaseIterator(values) { + + @Override + public boolean matches() throws IOException { + reader.reset(values.binaryValue()); + for (Component2DVisitor visitor : visitors) { + visitor.reset(); + reader.visit(visitor); + if (visitor.matches() == false) { + return false; + } + } + return true; + } + + @Override + public float matchCost() { + return 1000f; // TODO: what should it be? + } + }; + return new ConstantScoreScorer(this, boost, scoreMode, iterator); + } + + @Override + public boolean isCacheable(LeafReaderContext ctx) { + return DocValues.isCacheable(ctx, field); + } + }; + } +} diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/VectorGeoShapeWithDocValuesQueryProcessor.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/VectorGeoShapeWithDocValuesQueryProcessor.java new file mode 100644 index 0000000000000..7148e66cde553 --- /dev/null +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/VectorGeoShapeWithDocValuesQueryProcessor.java @@ -0,0 +1,77 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.xpack.spatial.index.query; + +import org.apache.lucene.document.LatLonShape; +import org.apache.lucene.geo.LatLonGeometry; +import org.apache.lucene.search.IndexOrDocValuesQuery; +import org.apache.lucene.search.MatchNoDocsQuery; +import org.apache.lucene.search.Query; +import org.elasticsearch.Version; +import org.elasticsearch.common.geo.GeoShapeUtils; +import org.elasticsearch.common.geo.ShapeRelation; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.Line; +import org.elasticsearch.geometry.MultiLine; +import org.elasticsearch.index.query.QueryShardException; +import org.elasticsearch.index.query.SearchExecutionContext; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + + +public class VectorGeoShapeWithDocValuesQueryProcessor { + + private static final List> WITHIN_UNSUPPORTED_GEOMETRIES = new ArrayList<>(); + static { + WITHIN_UNSUPPORTED_GEOMETRIES.add(Line.class); + WITHIN_UNSUPPORTED_GEOMETRIES.add(MultiLine.class); + } + + public Query geoShapeQuery(Geometry shape, String fieldName, ShapeRelation relation, + SearchExecutionContext context, boolean hasDocValues) { + // CONTAINS queries are not supported by VECTOR strategy for indices created before version 7.5.0 (Lucene 8.3.0) + if (relation == ShapeRelation.CONTAINS && context.indexVersionCreated().before(Version.V_7_5_0)) { + throw new QueryShardException(context, + ShapeRelation.CONTAINS + " query relation not supported for Field [" + fieldName + "]."); + } + final LatLonGeometry[] luceneGeometries = relation == ShapeRelation.WITHIN ? + GeoShapeUtils.toLuceneGeometry(fieldName, context, shape, WITHIN_UNSUPPORTED_GEOMETRIES) : + GeoShapeUtils.toLuceneGeometry(fieldName, context, shape, Collections.emptyList()); + + if (luceneGeometries.length == 0) { + return new MatchNoDocsQuery(); + } + Query query = LatLonShape.newGeometryQuery(fieldName, relation.getLuceneRelation(), luceneGeometries); + if (hasDocValues) { + final Query queryDocValues = new LatLonShapeDocValuesQuery(fieldName, relation.getLuceneRelation(), luceneGeometries); + query = new IndexOrDocValuesQuery(query, queryDocValues); + } + return query; + } +} + diff --git a/x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/index/query/GeoShapeWithDocValuesQueryBuilderTests.java b/x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/index/query/GeoShapeWithDocValuesQueryBuilderTests.java new file mode 100644 index 0000000000000..2cf5c12549433 --- /dev/null +++ b/x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/index/query/GeoShapeWithDocValuesQueryBuilderTests.java @@ -0,0 +1,66 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.spatial.index.query; + +import org.apache.lucene.search.ConstantScoreQuery; +import org.apache.lucene.search.IndexOrDocValuesQuery; +import org.apache.lucene.search.Query; +import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest; +import org.elasticsearch.common.Strings; +import org.elasticsearch.common.compress.CompressedXContent; +import org.elasticsearch.geo.GeometryTestUtils; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.index.mapper.MappedFieldType; +import org.elasticsearch.index.mapper.MapperService; +import org.elasticsearch.index.query.GeoShapeQueryBuilder; +import org.elasticsearch.index.query.SearchExecutionContext; +import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.test.AbstractQueryTestCase; +import org.elasticsearch.xpack.spatial.LocalStateSpatialPlugin; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Collection; + +import static org.hamcrest.Matchers.equalTo; + +public class GeoShapeWithDocValuesQueryBuilderTests extends AbstractQueryTestCase { + + @Override + protected Collection> getPlugins() { + return Arrays.asList(LocalStateSpatialPlugin.class); + } + + @Override + protected void initializeAdditionalMappings(MapperService mapperService) throws IOException { + if (randomBoolean()) { + mapperService.merge("_doc", new CompressedXContent(Strings.toString(PutMappingRequest.simpleMapping( + "test", "type=geo_shape"))), MapperService.MergeReason.MAPPING_UPDATE); + } else { + mapperService.merge("_doc", new CompressedXContent(Strings.toString(PutMappingRequest.simpleMapping( + "test", "type=geo_shape,doc_values=false"))), MapperService.MergeReason.MAPPING_UPDATE); + } + } + + @Override + protected GeoShapeQueryBuilder doCreateTestQueryBuilder() { + Geometry geometry = randomFrom( + GeometryTestUtils.randomPoint(false), + GeometryTestUtils.randomLine(false), + GeometryTestUtils.randomPolygon(false)); + return new GeoShapeQueryBuilder("test", geometry); + } + + @Override + protected void doAssertLuceneQuery(GeoShapeQueryBuilder queryBuilder, Query query, SearchExecutionContext context) { + assertThat(true, equalTo(query instanceof ConstantScoreQuery)); + Query geoShapeQuery = ((ConstantScoreQuery) query).getQuery(); + MappedFieldType fieldType = context.getFieldType("test"); + boolean IndexOrDocValuesQuery = fieldType.hasDocValues(); + assertThat(IndexOrDocValuesQuery, equalTo(geoShapeQuery instanceof IndexOrDocValuesQuery)); + } +} diff --git a/x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/index/query/LatLonShapeDocValuesQueryTests.java b/x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/index/query/LatLonShapeDocValuesQueryTests.java new file mode 100644 index 0000000000000..dd8f948344a4b --- /dev/null +++ b/x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/index/query/LatLonShapeDocValuesQueryTests.java @@ -0,0 +1,170 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ + +package org.elasticsearch.xpack.spatial.index.query; + +import org.apache.lucene.document.Document; +import org.apache.lucene.document.LatLonShape; +import org.apache.lucene.document.ShapeField; +import org.apache.lucene.geo.GeoTestUtil; +import org.apache.lucene.geo.LatLonGeometry; +import org.apache.lucene.geo.Point; +import org.apache.lucene.geo.Polygon; +import org.apache.lucene.geo.Rectangle; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.SerialMergeScheduler; +import org.apache.lucene.search.CheckHits; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryUtils; +import org.apache.lucene.store.Directory; +import org.elasticsearch.core.internal.io.IOUtils; +import org.elasticsearch.geo.GeometryTestUtils; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.index.mapper.GeoShapeIndexer; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.spatial.index.mapper.BinaryGeoShapeDocValuesField; + +import java.io.IOException; +import java.util.List; +import java.util.function.Function; + +public class LatLonShapeDocValuesQueryTests extends ESTestCase { + + private static final String FIELD_NAME = "field"; + + public void testEqualsAndHashcode() { + Polygon polygon = GeoTestUtil.nextPolygon(); + Query q1 = new LatLonShapeDocValuesQuery(FIELD_NAME,ShapeField.QueryRelation.INTERSECTS, polygon); + Query q2 = new LatLonShapeDocValuesQuery(FIELD_NAME,ShapeField.QueryRelation.INTERSECTS, polygon); + QueryUtils.checkEqual(q1, q2); + + Query q3 = new LatLonShapeDocValuesQuery(FIELD_NAME + "x",ShapeField.QueryRelation.INTERSECTS, polygon); + QueryUtils.checkUnequal(q1, q3); + + Rectangle rectangle = GeoTestUtil.nextBox(); + Query q4 = new LatLonShapeDocValuesQuery(FIELD_NAME,ShapeField.QueryRelation.INTERSECTS, rectangle); + QueryUtils.checkUnequal(q1, q4); + } + + public void testIndexSimpleShapes() throws Exception { + IndexWriterConfig iwc = newIndexWriterConfig(); + // Else seeds may not reproduce: + iwc.setMergeScheduler(new SerialMergeScheduler()); + // Else we can get O(N^2) merging: + iwc.setMaxBufferedDocs(10); + Directory dir = newDirectory(); + // RandomIndexWriter is too slow here: + IndexWriter w = new IndexWriter(dir, iwc); + final int numDocs = randomIntBetween(10, 1000); + GeoShapeIndexer indexer = new GeoShapeIndexer(true, FIELD_NAME); + for (int id = 0; id < numDocs; id++) { + Document doc = new Document(); + @SuppressWarnings("unchecked") Function geometryFunc = ESTestCase.randomFrom( + GeometryTestUtils::randomLine, + GeometryTestUtils::randomPoint, + GeometryTestUtils::randomPolygon + ); + Geometry geometry = geometryFunc.apply(false); + geometry = indexer.prepareForIndexing(geometry); + List fields = indexer.indexShape(null, geometry); + for (IndexableField field : fields) { + doc.add(field); + } + BinaryGeoShapeDocValuesField docVal = new BinaryGeoShapeDocValuesField(FIELD_NAME); + docVal.add(fields, geometry); + doc.add(docVal); + w.addDocument(doc); + } + + if (random().nextBoolean()) { + w.forceMerge(1); + } + final IndexReader r = DirectoryReader.open(w); + w.close(); + + IndexSearcher s = newSearcher(r); + for (int i = 0; i < 25; i++) { + LatLonGeometry[] geometries = randomLuceneQueryGeometries(); + for (ShapeField.QueryRelation relation : ShapeField.QueryRelation.values()) { + Query indexQuery = LatLonShape.newGeometryQuery(FIELD_NAME, relation, geometries); + Query docValQuery = new LatLonShapeDocValuesQuery(FIELD_NAME, relation, geometries); + assertQueries(s, indexQuery, docValQuery, numDocs); + } + } + IOUtils.close(r, dir); + } + + public void testIndexMultiShapes() throws Exception { + IndexWriterConfig iwc = newIndexWriterConfig(); + // Else seeds may not reproduce: + iwc.setMergeScheduler(new SerialMergeScheduler()); + // Else we can get O(N^2) merging: + iwc.setMaxBufferedDocs(10); + Directory dir = newDirectory(); + // RandomIndexWriter is too slow here: + IndexWriter w = new IndexWriter(dir, iwc); + final int numDocs = randomIntBetween(10, 100); + GeoShapeIndexer indexer = new GeoShapeIndexer(true, FIELD_NAME); + for (int id = 0; id < numDocs; id++) { + Document doc = new Document(); + Geometry geometry = GeometryTestUtils.randomGeometryWithoutCircle(randomIntBetween(1, 5), false); + geometry = indexer.prepareForIndexing(geometry); + List fields = indexer.indexShape(null, geometry); + for (IndexableField field : fields) { + doc.add(field); + } + BinaryGeoShapeDocValuesField docVal = new BinaryGeoShapeDocValuesField(FIELD_NAME); + docVal.add(fields, geometry); + doc.add(docVal); + w.addDocument(doc); + } + + if (random().nextBoolean()) { + w.forceMerge(1); + } + final IndexReader r = DirectoryReader.open(w); + w.close(); + + IndexSearcher s = newSearcher(r); + for (int i = 0; i < 25; i++) { + LatLonGeometry[] geometries = randomLuceneQueryGeometries(); + for (ShapeField.QueryRelation relation : ShapeField.QueryRelation.values()) { + Query indexQuery = LatLonShape.newGeometryQuery(FIELD_NAME, relation, geometries); + Query docValQuery = new LatLonShapeDocValuesQuery(FIELD_NAME, relation, geometries); + assertQueries(s, indexQuery, docValQuery, numDocs); + } + } + IOUtils.close(r, dir); + } + + private void assertQueries(IndexSearcher s, Query indexQuery, Query docValQuery, int numDocs) throws IOException { + assertEquals(s.count(indexQuery), s.count(docValQuery)); + CheckHits.checkEqual(docValQuery, s.search(indexQuery, numDocs).scoreDocs, s.search(docValQuery, numDocs).scoreDocs); + } + + private LatLonGeometry[] randomLuceneQueryGeometries() { + int numGeom = randomIntBetween(1, 3); + LatLonGeometry[] geometries = new LatLonGeometry[numGeom]; + for (int i = 0; i < numGeom; i++) { + geometries[i] = randomLuceneQueryGeometry(); + } + return geometries; + } + + private LatLonGeometry randomLuceneQueryGeometry() { + switch (randomInt(3)) { + case 0: return GeoTestUtil.nextPolygon(); + case 1: return GeoTestUtil.nextCircle(); + case 2: return new Point(GeoTestUtil.nextLatitude(), GeoTestUtil.nextLongitude()); + default: return GeoTestUtil.nextBox(); + } + } +}