diff --git a/server/src/main/java/org/elasticsearch/common/geo/CentroidCalculator.java b/server/src/main/java/org/elasticsearch/common/geo/CentroidCalculator.java new file mode 100644 index 0000000000000..c15c24d65ec40 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/common/geo/CentroidCalculator.java @@ -0,0 +1,91 @@ +/* + * 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.common.geo; + +/** + * This class keeps a running Kahan-sum of coordinates + * that are to be averaged in {@link GeometryTreeWriter} for use + * as the centroid of a shape. + */ +public class CentroidCalculator { + + private double compX; + private double compY; + private double sumX; + private double sumY; + private int count; + + public CentroidCalculator() { + this.sumX = 0.0; + this.compX = 0.0; + this.sumY = 0.0; + this.compY = 0.0; + this.count = 0; + } + + /** + * adds a single coordinate to the running sum and count of coordinates + * for centroid calculation + * + * @param x the x-coordinate of the point + * @param y the y-coordinate of the point + */ + public void addCoordinate(double x, double y) { + double correctedX = x - compX; + double newSumX = sumX + correctedX; + compX = (newSumX - sumX) - correctedX; + sumX = newSumX; + + double correctedY = y - compY; + double newSumY = sumY + correctedY; + compY = (newSumY - sumY) - correctedY; + sumY = newSumY; + + count += 1; + } + + /** + * Adjusts the existing calculator to add the running sum and count + * from another {@link CentroidCalculator}. This is used to keep + * a running count of points from different sub-shapes of a single + * geo-shape field + * + * @param otherCalculator the other centroid calculator to add from + */ + void addFrom(CentroidCalculator otherCalculator) { + addCoordinate(otherCalculator.sumX, otherCalculator.sumY); + // adjust count + count += otherCalculator.count - 1; + } + + /** + * @return the x-coordinate centroid + */ + public double getX() { + return sumX / count; + } + + /** + * @return the y-coordinate centroid + */ + public double getY() { + return sumY / count; + } +} diff --git a/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java b/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java index 2fb2f1ef230fe..3b11d00833cfc 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java +++ b/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java @@ -39,6 +39,7 @@ public class EdgeTreeWriter extends ShapeTreeWriter { private final Extent extent; private final int numShapes; + private final CentroidCalculator centroidCalculator; final Edge tree; @@ -46,12 +47,14 @@ public class EdgeTreeWriter extends ShapeTreeWriter { * @param x array of the x-coordinate of points. * @param y array of the y-coordinate of points. * @param coordinateEncoder class that encodes from real-valued x/y to serialized integer coordinate values. + * @param hasArea whether the tree represents a Polygon that has a defined area */ - EdgeTreeWriter(double[] x, double[] y, CoordinateEncoder coordinateEncoder) { - this(Collections.singletonList(x), Collections.singletonList(y), coordinateEncoder); + EdgeTreeWriter(double[] x, double[] y, CoordinateEncoder coordinateEncoder, boolean hasArea) { + this(Collections.singletonList(x), Collections.singletonList(y), coordinateEncoder, hasArea); } - EdgeTreeWriter(List x, List y, CoordinateEncoder coordinateEncoder) { + EdgeTreeWriter(List x, List y, CoordinateEncoder coordinateEncoder, boolean hasArea) { + this.centroidCalculator = new CentroidCalculator(); this.numShapes = x.size(); double top = Double.NEGATIVE_INFINITY; double bottom = Double.POSITIVE_INFINITY; @@ -59,6 +62,7 @@ public class EdgeTreeWriter extends ShapeTreeWriter { double negRight = Double.NEGATIVE_INFINITY; double posLeft = Double.POSITIVE_INFINITY; double posRight = Double.NEGATIVE_INFINITY; + List edges = new ArrayList<>(); for (int i = 0; i < y.size(); i++) { for (int j = 1; j < y.get(i).length; j++) { @@ -108,6 +112,12 @@ public class EdgeTreeWriter extends ShapeTreeWriter { if (x2 < 0 && x2 > negRight) { negRight = x2; } + + // calculate centroid + centroidCalculator.addCoordinate(x1, y1); + if (j == y.get(i).length - 1 && hasArea == false) { + centroidCalculator.addCoordinate(x2, y2); + } } } edges.sort(Edge::compareTo); @@ -127,6 +137,11 @@ public ShapeType getShapeType() { return numShapes > 1 ? ShapeType.MULTILINESTRING: ShapeType.LINESTRING; } + @Override + public CentroidCalculator getCentroidCalculator() { + return centroidCalculator; + } + @Override public void writeTo(StreamOutput out) throws IOException { extent.writeTo(out); diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java index 9f442d29cc41e..364ee63ce1ba4 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java +++ b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java @@ -34,14 +34,17 @@ */ public class GeometryTreeReader { + private final int extentOffset = 8; private final ByteBufferStreamInput input; + private final CoordinateEncoder coordinateEncoder; - public GeometryTreeReader(BytesRef bytesRef) { + public GeometryTreeReader(BytesRef bytesRef, CoordinateEncoder coordinateEncoder) { this.input = new ByteBufferStreamInput(ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length)); + this.coordinateEncoder = coordinateEncoder; } public Extent getExtent() throws IOException { - input.position(0); + input.position(extentOffset); Extent extent = input.readOptionalWriteable(Extent::new); if (extent != null) { return extent; @@ -52,8 +55,18 @@ public Extent getExtent() throws IOException { return reader.getExtent(); } - public boolean intersects(Extent extent) throws IOException { + public double getCentroidX() throws IOException { input.position(0); + return coordinateEncoder.decodeX(input.readInt()); + } + + public double getCentroidY() throws IOException { + input.position(4); + return coordinateEncoder.decodeY(input.readInt()); + } + + public boolean intersects(Extent extent) throws IOException { + input.position(extentOffset); boolean hasExtent = input.readBoolean(); if (hasExtent) { Optional extentCheck = EdgeTreeReader.checkExtent(new Extent(input), extent); diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java index a6c1dfb3a6300..61bda1d8c44ed 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java +++ b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java @@ -46,8 +46,12 @@ public class GeometryTreeWriter implements Writeable { private final GeometryTreeBuilder builder; + private final CoordinateEncoder coordinateEncoder; + private CentroidCalculator centroidCalculator; public GeometryTreeWriter(Geometry geometry, CoordinateEncoder coordinateEncoder) { + this.coordinateEncoder = coordinateEncoder; + this.centroidCalculator = new CentroidCalculator(); builder = new GeometryTreeBuilder(coordinateEncoder); geometry.visit(builder); } @@ -62,6 +66,8 @@ public void writeTo(StreamOutput out) throws IOException { // contains multiple sub-shapes boolean prependExtent = builder.shapeWriters.size() > 1; Extent extent = null; + out.writeInt(coordinateEncoder.encodeX(centroidCalculator.getX())); + out.writeInt(coordinateEncoder.encodeY(centroidCalculator.getY())); if (prependExtent) { extent = new Extent(builder.top, builder.bottom, builder.negLeft, builder.negRight, builder.posLeft, builder.posRight); } @@ -99,6 +105,7 @@ private void addWriter(ShapeTreeWriter writer) { posLeft = Math.min(posLeft, extent.posLeft); posRight = Math.max(posRight, extent.posRight); shapeWriters.add(writer); + centroidCalculator.addFrom(writer.getCentroidCalculator()); } @Override @@ -111,7 +118,7 @@ public Void visit(GeometryCollection collection) { @Override public Void visit(Line line) { - addWriter(new EdgeTreeWriter(line.getLons(), line.getLats(), coordinateEncoder)); + addWriter(new EdgeTreeWriter(line.getLons(), line.getLats(), coordinateEncoder, false)); return null; } @@ -124,7 +131,7 @@ public Void visit(MultiLine multiLine) { x.add(line.getLons()); y.add(line.getLats()); } - addWriter(new EdgeTreeWriter(x, y, coordinateEncoder)); + addWriter(new EdgeTreeWriter(x, y, coordinateEncoder, false)); return null; } @@ -181,7 +188,7 @@ public Void visit(MultiPoint multiPoint) { @Override public Void visit(LinearRing ring) { - throw new IllegalArgumentException("invalid shape type found [Circle]"); + throw new IllegalArgumentException("invalid shape type found [LinearRing]"); } @Override diff --git a/server/src/main/java/org/elasticsearch/common/geo/Point2DWriter.java b/server/src/main/java/org/elasticsearch/common/geo/Point2DWriter.java index f59a806f958ab..bd2a250925b78 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/Point2DWriter.java +++ b/server/src/main/java/org/elasticsearch/common/geo/Point2DWriter.java @@ -36,10 +36,12 @@ public class Point2DWriter extends ShapeTreeWriter { // size of a leaf node where searches are done sequentially. static final int LEAF_SIZE = 64; private final CoordinateEncoder coordinateEncoder; + private final CentroidCalculator centroidCalculator; Point2DWriter(double[] x, double[] y, CoordinateEncoder coordinateEncoder) { assert x.length == y.length; this.coordinateEncoder = coordinateEncoder; + this.centroidCalculator = new CentroidCalculator(); double top = Double.NEGATIVE_INFINITY; double bottom = Double.POSITIVE_INFINITY; double negLeft = Double.POSITIVE_INFINITY; @@ -47,6 +49,7 @@ public class Point2DWriter extends ShapeTreeWriter { double posLeft = Double.POSITIVE_INFINITY; double posRight = Double.NEGATIVE_INFINITY; coords = new double[x.length * K]; + for (int i = 0; i < x.length; i++) { double xi = x[i]; double yi = y[i]; @@ -66,6 +69,8 @@ public class Point2DWriter extends ShapeTreeWriter { } coords[2 * i] = xi; coords[2 * i + 1] = yi; + + centroidCalculator.addCoordinate(xi, yi); } sort(0, x.length - 1, 0); this.extent = new Extent(coordinateEncoder.encodeY(top), coordinateEncoder.encodeY(bottom), coordinateEncoder.encodeX(negLeft), @@ -76,6 +81,8 @@ public class Point2DWriter extends ShapeTreeWriter { this.coordinateEncoder = coordinateEncoder; coords = new double[] {x, y}; this.extent = Extent.fromPoint(coordinateEncoder.encodeX(x), coordinateEncoder.encodeY(y)); + this.centroidCalculator = new CentroidCalculator(); + centroidCalculator.addCoordinate(x, y); } @Override @@ -88,6 +95,11 @@ public ShapeType getShapeType() { return ShapeType.MULTIPOINT; } + @Override + public CentroidCalculator getCentroidCalculator() { + return centroidCalculator; + } + @Override public void writeTo(StreamOutput out) throws IOException { int numPoints = coords.length >> 1; diff --git a/server/src/main/java/org/elasticsearch/common/geo/PolygonTreeWriter.java b/server/src/main/java/org/elasticsearch/common/geo/PolygonTreeWriter.java index b48ce0e1cb31b..d114ae40f84c7 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/PolygonTreeWriter.java +++ b/server/src/main/java/org/elasticsearch/common/geo/PolygonTreeWriter.java @@ -35,8 +35,8 @@ public class PolygonTreeWriter extends ShapeTreeWriter { private final EdgeTreeWriter holes; public PolygonTreeWriter(double[] x, double[] y, List holesX, List holesY, CoordinateEncoder coordinateEncoder) { - outerShell = new EdgeTreeWriter(x, y, coordinateEncoder); - holes = holesX.isEmpty() ? null : new EdgeTreeWriter(holesX, holesY, coordinateEncoder); + outerShell = new EdgeTreeWriter(x, y, coordinateEncoder, true); + holes = holesX.isEmpty() ? null : new EdgeTreeWriter(holesX, holesY, coordinateEncoder, true); } public Extent getExtent() { @@ -47,6 +47,11 @@ public ShapeType getShapeType() { return ShapeType.POLYGON; } + @Override + public CentroidCalculator getCentroidCalculator() { + return outerShell.getCentroidCalculator(); + } + @Override public void writeTo(StreamOutput out) throws IOException { // calculate size of outerShell's tree to make it easy to jump to the holes tree quickly when querying diff --git a/server/src/main/java/org/elasticsearch/common/geo/ShapeTreeWriter.java b/server/src/main/java/org/elasticsearch/common/geo/ShapeTreeWriter.java index 11f555b7176d5..35eaec5fb02f0 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/ShapeTreeWriter.java +++ b/server/src/main/java/org/elasticsearch/common/geo/ShapeTreeWriter.java @@ -29,4 +29,6 @@ public abstract class ShapeTreeWriter implements Writeable { public abstract Extent getExtent(); public abstract ShapeType getShapeType(); + + public abstract CentroidCalculator getCentroidCalculator(); } diff --git a/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java b/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java index 6486c74c7fb2a..c8dc4f7cb44b5 100644 --- a/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java +++ b/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java @@ -128,14 +128,28 @@ public BoundingBox boundingBox() { return new BoundingBox(extent, GeoShapeCoordinateEncoder.INSTANCE); } + /** + * @return the latitude of the centroid of the shape + */ @Override public double lat() { - throw new UnsupportedOperationException("centroid of GeoShape is not defined"); + try { + return reader.getCentroidY(); + } catch (IOException e) { + throw new IllegalStateException("unable to read centroid of shape", e); + } } + /** + * @return the longitude of the centroid of the shape + */ @Override public double lon() { - throw new UnsupportedOperationException("centroid of GeoShape is not defined"); + try { + return reader.getCentroidX(); + } catch (IOException e) { + throw new IllegalStateException("unable to read centroid of shape", e); + } } public static GeoShapeValue missing(String missing) { diff --git a/server/src/main/java/org/elasticsearch/index/fielddata/plain/LatLonShapeDVAtomicFieldData.java b/server/src/main/java/org/elasticsearch/index/fielddata/plain/LatLonShapeDVAtomicFieldData.java index ec03a3959d482..2e0480263174c 100644 --- a/server/src/main/java/org/elasticsearch/index/fielddata/plain/LatLonShapeDVAtomicFieldData.java +++ b/server/src/main/java/org/elasticsearch/index/fielddata/plain/LatLonShapeDVAtomicFieldData.java @@ -23,6 +23,7 @@ import org.apache.lucene.index.LeafReader; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.geo.GeoShapeCoordinateEncoder; import org.elasticsearch.common.geo.GeometryTreeReader; import org.elasticsearch.index.fielddata.MultiGeoValues; @@ -74,7 +75,7 @@ public int docValueCount() { @Override public GeoValue nextValue() throws IOException { final BytesRef encoded = binaryValues.binaryValue(); - return new GeoShapeValue(new GeometryTreeReader(encoded)); + return new GeoShapeValue(new GeometryTreeReader(encoded, GeoShapeCoordinateEncoder.INSTANCE)); } }; } catch (IOException e) { diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregationBuilder.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregationBuilder.java index af2d1a600243a..6f3d360abb0f8 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregationBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregationBuilder.java @@ -39,13 +39,13 @@ import java.util.Map; public class GeoCentroidAggregationBuilder - extends ValuesSourceAggregationBuilder.LeafOnly { + extends ValuesSourceAggregationBuilder.LeafOnly { public static final String NAME = "geo_centroid"; private static final ObjectParser PARSER; static { PARSER = new ObjectParser<>(GeoCentroidAggregationBuilder.NAME); - ValuesSourceParserHelper.declareGeoPointFields(PARSER, true, false); + ValuesSourceParserHelper.declareGeoFields(PARSER, true, false); } public static AggregationBuilder parse(String aggregationName, XContentParser parser) throws IOException { @@ -78,7 +78,7 @@ protected void innerWriteTo(StreamOutput out) { } @Override - protected GeoCentroidAggregatorFactory innerBuild(SearchContext context, ValuesSourceConfig config, + protected GeoCentroidAggregatorFactory innerBuild(SearchContext context, ValuesSourceConfig config, AggregatorFactory parent, Builder subFactoriesBuilder) throws IOException { return new GeoCentroidAggregatorFactory(name, config, context, parent, subFactoriesBuilder, metaData); } diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregator.java index 975356c7d6fe3..c91ebb92e8714 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregator.java @@ -42,12 +42,12 @@ * A geo metric aggregator that computes a geo-centroid from a {@code geo_point} type field */ final class GeoCentroidAggregator extends MetricsAggregator { - private final ValuesSource.GeoPoint valuesSource; + private final ValuesSource.Geo valuesSource; private DoubleArray lonSum, lonCompensations, latSum, latCompensations; private LongArray counts; GeoCentroidAggregator(String name, SearchContext context, Aggregator parent, - ValuesSource.GeoPoint valuesSource, List pipelineAggregators, + ValuesSource.Geo valuesSource, List pipelineAggregators, Map metaData) throws IOException { super(name, context, parent, pipelineAggregators, metaData); this.valuesSource = valuesSource; @@ -89,6 +89,9 @@ public void collect(int doc, long bucket) throws IOException { double compensationLon = lonCompensations.get(bucket); // update the sum + // + // this calculates the centroid of centroid of shapes when + // executing against geo-shape fields. for (int i = 0; i < valueCount; ++i) { MultiGeoValues.GeoValue value = values.nextValue(); //latitude diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorFactory.java index b12ce921b7d52..3dd694fc21b7f 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorFactory.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorFactory.java @@ -32,9 +32,9 @@ import java.util.List; import java.util.Map; -class GeoCentroidAggregatorFactory extends ValuesSourceAggregatorFactory { +class GeoCentroidAggregatorFactory extends ValuesSourceAggregatorFactory { - GeoCentroidAggregatorFactory(String name, ValuesSourceConfig config, + GeoCentroidAggregatorFactory(String name, ValuesSourceConfig config, SearchContext context, AggregatorFactory parent, AggregatorFactories.Builder subFactoriesBuilder, Map metaData) throws IOException { super(name, config, context, parent, subFactoriesBuilder, metaData); @@ -47,7 +47,7 @@ protected Aggregator createUnmapped(Aggregator parent, } @Override - protected Aggregator doCreateInternal(ValuesSource.GeoPoint valuesSource, Aggregator parent, + protected Aggregator doCreateInternal(ValuesSource.Geo valuesSource, Aggregator parent, boolean collectsFromSingleBucket, List pipelineAggregators, Map metaData) throws IOException { return new GeoCentroidAggregator(name, context, parent, valuesSource, pipelineAggregators, metaData); diff --git a/server/src/test/java/org/elasticsearch/common/geo/CentroidCalculatorTests.java b/server/src/test/java/org/elasticsearch/common/geo/CentroidCalculatorTests.java new file mode 100644 index 0000000000000..43c5e4a02134b --- /dev/null +++ b/server/src/test/java/org/elasticsearch/common/geo/CentroidCalculatorTests.java @@ -0,0 +1,44 @@ +/* + * 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.common.geo; + +import org.elasticsearch.test.ESTestCase; + +import static org.hamcrest.Matchers.equalTo; + +public class CentroidCalculatorTests extends ESTestCase { + + public void test() { + double[] x = new double[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }; + double[] y = new double[] { 10, 20, 30, 40, 50, 60, 70, 80, 90, 100 }; + double[] xRunningAvg = new double[] { 1, 1.5, 2.0, 2.5, 3, 3.5, 4, 4.5, 5, 5.5 }; + double[] yRunningAvg = new double[] { 10, 15, 20, 25, 30, 35, 40, 45, 50, 55 }; + CentroidCalculator calculator = new CentroidCalculator(); + for (int i = 0; i < 10; i++) { + calculator.addCoordinate(x[i], y[i]); + assertThat(calculator.getX(), equalTo(xRunningAvg[i])); + assertThat(calculator.getY(), equalTo(yRunningAvg[i])); + } + CentroidCalculator otherCalculator = new CentroidCalculator(); + otherCalculator.addCoordinate(0.0, 0.0); + calculator.addFrom(otherCalculator); + assertThat(calculator.getX(), equalTo(5.0)); + assertThat(calculator.getY(), equalTo(50.0)); + } +} diff --git a/server/src/test/java/org/elasticsearch/common/geo/EdgeTreeTests.java b/server/src/test/java/org/elasticsearch/common/geo/EdgeTreeTests.java index 802252bb63674..4bdeba2e55612 100644 --- a/server/src/test/java/org/elasticsearch/common/geo/EdgeTreeTests.java +++ b/server/src/test/java/org/elasticsearch/common/geo/EdgeTreeTests.java @@ -47,11 +47,15 @@ public void testRectangleShape() throws IOException { int maxY = randomIntBetween(minY + 10, 180); double[] x = new double[]{minX, maxX, maxX, minX, minX}; double[] y = new double[]{minY, minY, maxY, maxY, minY}; - EdgeTreeWriter writer = new EdgeTreeWriter(x, y, TestCoordinateEncoder.INSTANCE); + EdgeTreeWriter writer = new EdgeTreeWriter(x, y, TestCoordinateEncoder.INSTANCE, true); BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - EdgeTreeReader reader = new EdgeTreeReader(new ByteBufferStreamInput(ByteBuffer.wrap(output.bytes().toBytesRef().bytes)), true); + EdgeTreeReader reader = new EdgeTreeReader( + new ByteBufferStreamInput(ByteBuffer.wrap(output.bytes().toBytesRef().bytes)), true); + + assertThat(writer.getCentroidCalculator().getX(), equalTo((minX + maxX)/2.0)); + assertThat(writer.getCentroidCalculator().getY(), equalTo((minY + maxY)/2.0)); // box-query touches bottom-left corner assertTrue(reader.intersects(Extent.fromPoints(minX - randomIntBetween(1, 180), minY - randomIntBetween(1, 180), minX, minY))); @@ -119,7 +123,7 @@ public void testSimplePolygon() throws IOException { double[] x = testPolygon.getPolygon().getLons(); double[] y = testPolygon.getPolygon().getLats(); - EdgeTreeWriter writer = new EdgeTreeWriter(x, y, TestCoordinateEncoder.INSTANCE); + EdgeTreeWriter writer = new EdgeTreeWriter(x, y, TestCoordinateEncoder.INSTANCE, true); BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); @@ -155,7 +159,7 @@ public void testPacMan() throws Exception { int yMax = 1;//5; // test cell crossing poly - EdgeTreeWriter writer = new EdgeTreeWriter(px, py, TestCoordinateEncoder.INSTANCE); + EdgeTreeWriter writer = new EdgeTreeWriter(px, py, TestCoordinateEncoder.INSTANCE, true); BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); @@ -165,10 +169,10 @@ public void testPacMan() throws Exception { public void testGetShapeType() { double[] pointCoord = new double[] { 0 }; - assertThat(new EdgeTreeWriter(pointCoord, pointCoord, TestCoordinateEncoder.INSTANCE).getShapeType(), + assertThat(new EdgeTreeWriter(pointCoord, pointCoord, TestCoordinateEncoder.INSTANCE, false).getShapeType(), equalTo(ShapeType.LINESTRING)); assertThat(new EdgeTreeWriter(List.of(pointCoord, pointCoord), List.of(pointCoord, pointCoord), - TestCoordinateEncoder.INSTANCE).getShapeType(), + TestCoordinateEncoder.INSTANCE, false).getShapeType(), equalTo(ShapeType.MULTILINESTRING)); } } diff --git a/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java b/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java index 0387848f09a81..d8abe1aaadd69 100644 --- a/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java +++ b/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java @@ -52,9 +52,12 @@ public void testRectangleShape() throws IOException { BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef()); + GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef(), TestCoordinateEncoder.INSTANCE); assertThat(Extent.fromPoints(minX, minY, maxX, maxY), equalTo(reader.getExtent())); + // encoder loses precision when casting to integer, so centroid is calculated using integer division here + assertThat(reader.getCentroidX(), equalTo((double) ((minX + maxX)/2))); + assertThat(reader.getCentroidY(), equalTo((double) ((minY + maxY)/2))); // box-query touches bottom-left corner assertTrue(reader.intersects(Extent.fromPoints(minX - randomIntBetween(1, 10), minY - randomIntBetween(1, 10), minX, minY))); @@ -105,7 +108,7 @@ public void testPacManPolygon() throws Exception { BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef()); + GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef(), TestCoordinateEncoder.INSTANCE); assertTrue(reader.intersects(Extent.fromPoints(2, -1, 11, 1))); assertTrue(reader.intersects(Extent.fromPoints(-12, -12, 12, 12))); assertTrue(reader.intersects(Extent.fromPoints(-2, -1, 2, 0))); @@ -121,7 +124,7 @@ public void testPolygonWithHole() throws Exception { BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef()); + GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef(), null); assertFalse(reader.intersects(Extent.fromPoints(6, -6, 6, -6))); // in the hole assertTrue(reader.intersects(Extent.fromPoints(25, -25, 25, -25))); // on the mainland @@ -144,7 +147,7 @@ public void testCombPolygon() throws Exception { BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef()); + GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef(), TestCoordinateEncoder.INSTANCE); assertTrue(reader.intersects(Extent.fromPoints(5, 10, 5, 10))); assertFalse(reader.intersects(Extent.fromPoints(15, 10, 15, 10))); assertFalse(reader.intersects(Extent.fromPoints(25, 10, 25, 10))); @@ -160,7 +163,7 @@ public void testPacManClosedLineString() throws Exception { BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef()); + GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef(), TestCoordinateEncoder.INSTANCE); assertTrue(reader.intersects(Extent.fromPoints(2, -1, 11, 1))); assertTrue(reader.intersects(Extent.fromPoints(-12, -12, 12, 12))); assertTrue(reader.intersects(Extent.fromPoints(-2, -1, 2, 0))); @@ -177,7 +180,7 @@ public void testPacManLineString() throws Exception { BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef()); + GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef(), TestCoordinateEncoder.INSTANCE); assertTrue(reader.intersects(Extent.fromPoints(2, -1, 11, 1))); assertTrue(reader.intersects(Extent.fromPoints(-12, -12, 12, 12))); assertTrue(reader.intersects(Extent.fromPoints(-2, -1, 2, 0))); @@ -211,7 +214,7 @@ public void testPacManPoints() throws Exception { BytesStreamOutput output = new BytesStreamOutput(); writer.writeTo(output); output.close(); - GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef()); + GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef(), TestCoordinateEncoder.INSTANCE); assertTrue(reader.intersects(Extent.fromPoints(xMin, yMin, xMax, yMax))); } } diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/metrics/AbstractGeoTestCase.java b/server/src/test/java/org/elasticsearch/search/aggregations/metrics/AbstractGeoTestCase.java index 5d6c10079d145..3883eeeab9c49 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/metrics/AbstractGeoTestCase.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/metrics/AbstractGeoTestCase.java @@ -68,7 +68,7 @@ public abstract class AbstractGeoTestCase extends ESIntegTestCase { protected static int numUniqueGeoPoints; protected static GeoPoint[] singleValues, multiValues; protected static GeoPoint singleTopLeft, singleBottomRight, multiTopLeft, multiBottomRight, - singleCentroid, multiCentroid, unmappedCentroid; + singleCentroid, singleShapeCentroid, multiCentroid, unmappedCentroid; protected static ObjectIntMap expectedDocCountsForGeoHash = null; protected static ObjectObjectMap expectedCentroidsForGeoHash = null; protected static final double GEOHASH_TOLERANCE = 1E-5D; @@ -85,6 +85,7 @@ public void setupSuiteScopeCluster() throws Exception { multiTopLeft = new GeoPoint(Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY); multiBottomRight = new GeoPoint(Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY); singleCentroid = new GeoPoint(0, 0); + singleShapeCentroid = new GeoPoint(9.4, 34.4); multiCentroid = new GeoPoint(0, 0); unmappedCentroid = new GeoPoint(0, 0); diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorTests.java index 4acae1764953a..83e3887e323f8 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorTests.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorTests.java @@ -25,14 +25,32 @@ import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.store.Directory; +import org.elasticsearch.common.geo.CentroidCalculator; import org.elasticsearch.common.geo.GeoPoint; +import org.elasticsearch.geo.GeometryTestUtils; +import org.elasticsearch.geometry.Circle; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.GeometryCollection; +import org.elasticsearch.geometry.GeometryVisitor; +import org.elasticsearch.geometry.Line; +import org.elasticsearch.geometry.LinearRing; +import org.elasticsearch.geometry.MultiLine; +import org.elasticsearch.geometry.MultiPoint; +import org.elasticsearch.geometry.MultiPolygon; +import org.elasticsearch.geometry.Point; +import org.elasticsearch.geometry.Polygon; +import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.index.mapper.BinaryGeoShapeDocValuesField; import org.elasticsearch.index.mapper.GeoPointFieldMapper; +import org.elasticsearch.index.mapper.GeoShapeFieldMapper; import org.elasticsearch.index.mapper.MappedFieldType; import org.elasticsearch.search.aggregations.AggregatorTestCase; import org.elasticsearch.search.aggregations.support.AggregationInspectionHelper; +import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.geo.RandomGeoGenerator; import java.io.IOException; +import java.util.function.Function; public class GeoCentroidAggregatorTests extends AggregatorTestCase { @@ -84,7 +102,7 @@ public void testUnmapped() throws Exception { } } - public void testSingleValuedField() throws Exception { + public void testSingleValuedGeoPointField() throws Exception { int numDocs = scaledRandomIntBetween(64, 256); int numUniqueGeoPoints = randomIntBetween(1, numDocs); try (Directory dir = newDirectory(); @@ -103,11 +121,11 @@ public void testSingleValuedField() throws Exception { expectedCentroid = expectedCentroid.reset(expectedCentroid.lat() + (singleVal.lat() - expectedCentroid.lat()) / (i + 1), expectedCentroid.lon() + (singleVal.lon() - expectedCentroid.lon()) / (i + 1)); } - assertCentroid(w, expectedCentroid); + assertCentroid(w, expectedCentroid, new GeoPointFieldMapper.GeoPointFieldType()); } } - public void testMultiValuedField() throws Exception { + public void testMultiValuedGeoPointField() throws Exception { int numDocs = scaledRandomIntBetween(64, 256); int numUniqueGeoPoints = randomIntBetween(1, numDocs); try (Directory dir = newDirectory(); @@ -131,12 +149,115 @@ public void testMultiValuedField() throws Exception { expectedCentroid = expectedCentroid.reset(expectedCentroid.lat() + (newMVLat - expectedCentroid.lat()) / (i + 1), expectedCentroid.lon() + (newMVLon - expectedCentroid.lon()) / (i + 1)); } - assertCentroid(w, expectedCentroid); + assertCentroid(w, expectedCentroid, new GeoPointFieldMapper.GeoPointFieldType()); } } - private void assertCentroid(RandomIndexWriter w, GeoPoint expectedCentroid) throws IOException { - MappedFieldType fieldType = new GeoPointFieldMapper.GeoPointFieldType(); + @SuppressWarnings("unchecked") + public void testGeoShapeField() throws Exception { + int numDocs = scaledRandomIntBetween(64, 256); + Function geometryGenerator = ESTestCase.randomFrom( + GeometryTestUtils::randomLine, + GeometryTestUtils::randomPoint, + GeometryTestUtils::randomPolygon, + GeometryTestUtils::randomMultiLine, + GeometryTestUtils::randomMultiPoint, + (hasAlt) -> GeometryTestUtils.randomRectangle(), + GeometryTestUtils::randomMultiPolygon + ); + try (Directory dir = newDirectory(); + RandomIndexWriter w = new RandomIndexWriter(random(), dir)) { + GeoPoint expectedCentroid = new GeoPoint(0, 0); + CentroidCalculator centroidOfCentroidsCalculator = new CentroidCalculator(); + for (int i = 0; i < numDocs; i++) { + CentroidCalculator calculator = new CentroidCalculator(); + Document document = new Document(); + Geometry geometry = geometryGenerator.apply(false); + geometry.visit(new GeometryVisitor() { + @Override + public Void visit(Circle circle) throws Exception { + calculator.addCoordinate(circle.getX(), circle.getY()); + return null; + } + + @Override + public Void visit(GeometryCollection collection) throws Exception { + for (Geometry shape : collection) { + shape.visit(this); + } + return null; + } + + @Override + public Void visit(Line line) throws Exception { + for (int i = 0; i < line.length(); i++) { + calculator.addCoordinate(line.getX(i), line.getY(i)); + } + return null; + } + + @Override + public Void visit(LinearRing ring) throws Exception { + for (int i = 0; i < ring.length() - 1; i++) { + calculator.addCoordinate(ring.getX(i), ring.getY(i)); + } + return null; + } + + @Override + public Void visit(MultiLine multiLine) throws Exception { + for (Line line : multiLine) { + visit(line); + } + return null; + } + + @Override + public Void visit(MultiPoint multiPoint) throws Exception { + for (Point point : multiPoint) { + visit(point); + } + return null; + } + + @Override + public Void visit(MultiPolygon multiPolygon) throws Exception { + for (Polygon polygon : multiPolygon) { + visit(polygon); + } + return null; + } + + @Override + public Void visit(Point point) throws Exception { + calculator.addCoordinate(point.getX(), point.getY()); + return null; + } + + @Override + public Void visit(Polygon polygon) throws Exception { + return visit(polygon.getPolygon()); + } + + @Override + public Void visit(Rectangle rectangle) throws Exception { + calculator.addCoordinate(rectangle.getMinX(), rectangle.getMinY()); + calculator.addCoordinate(rectangle.getMinX(), rectangle.getMaxY()); + calculator.addCoordinate(rectangle.getMaxX(), rectangle.getMinY()); + calculator.addCoordinate(rectangle.getMaxX(), rectangle.getMaxY()); + return null; + } + }); + document.add(new BinaryGeoShapeDocValuesField("field", geometry)); + w.addDocument(document); + centroidOfCentroidsCalculator.addCoordinate(calculator.getX(), calculator.getY()); + } + expectedCentroid.reset(centroidOfCentroidsCalculator.getY(), centroidOfCentroidsCalculator.getX()); + assertCentroid(w, expectedCentroid, new GeoShapeFieldMapper.GeoShapeFieldType()); + } + } + + private void assertCentroid(RandomIndexWriter w, GeoPoint expectedCentroid, MappedFieldType fieldType) throws IOException { fieldType.setHasDocValues(true); fieldType.setName("field"); GeoCentroidAggregationBuilder aggBuilder = new GeoCentroidAggregationBuilder("my_agg") diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidIT.java b/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidIT.java index d941ce27ede2f..576c4a34e347c 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidIT.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidIT.java @@ -106,6 +106,22 @@ public void testSingleValuedField() throws Exception { assertEquals(numDocs, geoCentroid.count()); } + public void testShapeField() throws Exception { + SearchResponse response = client().prepareSearch(DATELINE_IDX_NAME) + .setQuery(matchAllQuery()) + .addAggregation(geoCentroid(aggName).field(SINGLE_VALUED_GEOSHAPE_FIELD_NAME)) + .get(); + assertSearchResponse(response); + + GeoCentroid geoCentroid = response.getAggregations().get(aggName); + assertThat(geoCentroid, notNullValue()); + assertThat(geoCentroid.getName(), equalTo(aggName)); + GeoPoint centroid = geoCentroid.centroid(); + assertThat(centroid.lat(), closeTo(singleShapeCentroid.lat(), GEOHASH_TOLERANCE)); + assertThat(centroid.lon(), closeTo(singleShapeCentroid.lon(), GEOHASH_TOLERANCE)); + assertEquals(5, geoCentroid.count()); + } + public void testSingleValueFieldGetProperty() throws Exception { SearchResponse response = client().prepareSearch(IDX_NAME) .setQuery(matchAllQuery()) @@ -123,7 +139,7 @@ public void testSingleValueFieldGetProperty() throws Exception { GeoCentroid geoCentroid = global.getAggregations().get(aggName); assertThat(geoCentroid, notNullValue()); assertThat(geoCentroid.getName(), equalTo(aggName)); - assertThat((GeoCentroid) ((InternalAggregation)global).getProperty(aggName), sameInstance(geoCentroid)); + assertThat(((InternalAggregation)global).getProperty(aggName), sameInstance(geoCentroid)); GeoPoint centroid = geoCentroid.centroid(); assertThat(centroid.lat(), closeTo(singleCentroid.lat(), GEOHASH_TOLERANCE)); assertThat(centroid.lon(), closeTo(singleCentroid.lon(), GEOHASH_TOLERANCE));