Add geo_shape support to geo_centroid aggregation

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;
+    }
+}

server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java

@@ -39,26 +39,30 @@ public class EdgeTreeWriter extends ShapeTreeWriter {
 
     private final Extent extent;
     private final int numShapes;
+    private final CentroidCalculator centroidCalculator;
     final Edge tree;
 
 
     /**
      * @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<double[]> x, List<double[]> y, CoordinateEncoder coordinateEncoder) {
+    EdgeTreeWriter(List<double[]> x, List<double[]> y, CoordinateEncoder coordinateEncoder, boolean hasArea) {
+        this.centroidCalculator = new CentroidCalculator();
         this.numShapes = x.size();
         double top = Double.NEGATIVE_INFINITY;
         double bottom = Double.POSITIVE_INFINITY;
         double negLeft = Double.POSITIVE_INFINITY;
         double negRight = Double.NEGATIVE_INFINITY;
         double posLeft = Double.POSITIVE_INFINITY;
         double posRight = Double.NEGATIVE_INFINITY;
+
         List<Edge> 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);

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<Boolean> extentCheck = EdgeTreeReader.checkExtent(new Extent(input), extent);

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

server/src/main/java/org/elasticsearch/common/geo/Point2DWriter.java

@@ -36,17 +36,20 @@ 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;
         double negRight = Double.NEGATIVE_INFINITY;
         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;

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<double[]> holesX, List<double[]> 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

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

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) {

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) {