Implement weighted geo_shape centroid support

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

@@ -38,20 +38,19 @@
  * as the centroid of a shape.
  */
 public class CentroidCalculator {
-
     private double compX;
     private double compY;
     private double sumX;
     private double sumY;
-    private int count;
+    private double sumWeight;
     private DimensionalShapeType dimensionalShapeType;
 
     public CentroidCalculator(Geometry geometry) {
         this.sumX = 0.0;
         this.compX = 0.0;
         this.sumY = 0.0;
         this.compY = 0.0;
-        this.count = 0;
+        this.sumWeight = 0.0;
         CentroidCalculatorVisitor visitor = new CentroidCalculatorVisitor(this);
         geometry.visit(visitor);
         this.dimensionalShapeType = DimensionalShapeType.forGeometry(geometry);
@@ -60,22 +59,22 @@ public CentroidCalculator(Geometry geometry) {
     /**
      * adds a single coordinate to the running sum and count of coordinates
      * for centroid calculation
-     *
-     * @param x the x-coordinate of the point
+     *  @param x the x-coordinate of the point
      * @param y the y-coordinate of the point
+     * @param weight the associated weight of the coordinate
      */
-    private void addCoordinate(double x, double y) {
-        double correctedX = x - compX;
+    private void addCoordinate(double x, double y, double weight) {
+        double correctedX = weight * x - compX;
         double newSumX = sumX + correctedX;
         compX = (newSumX - sumX) - correctedX;
         sumX = newSumX;
 
-        double correctedY = y - compY;
+        double correctedY = weight * y - compY;
         double newSumY = sumY + correctedY;
         compY = (newSumY - sumY) - correctedY;
         sumY = newSumY;
 
-        count += 1;
+        sumWeight += weight;
     }
 
     /**
@@ -87,26 +86,45 @@ private void addCoordinate(double x, double y) {
      * @param otherCalculator the other centroid calculator to add from
      */
     public void addFrom(CentroidCalculator otherCalculator) {
-        addCoordinate(otherCalculator.sumX, otherCalculator.sumY);
-        // adjust count
-        count += otherCalculator.count - 1;
-        dimensionalShapeType = DimensionalShapeType.max(dimensionalShapeType, otherCalculator.dimensionalShapeType);
+        int compared = DimensionalShapeType.COMPARATOR.compare(dimensionalShapeType, otherCalculator.dimensionalShapeType);
+        if (compared < 0) {
+            sumWeight = otherCalculator.sumWeight;
+            dimensionalShapeType = otherCalculator.dimensionalShapeType;
+            sumX = otherCalculator.sumX;
+            sumY = otherCalculator.sumY;
+            compX = otherCalculator.compX;
+            compY = otherCalculator.compY;
+        } else if (compared == 0) {
+            addCoordinate(otherCalculator.sumX, otherCalculator.sumY, otherCalculator.sumWeight);
+        } // else (compared > 0) do not modify centroid calculation since otherCalculator is of lower dimension than this calculator
     }
 
     /**
      * @return the x-coordinate centroid
      */
     public double getX() {
-        return sumX / count;
+        // normalization required due to floating point precision errors
+        return GeoUtils.normalizeLon(sumX / sumWeight);
     }
 
     /**
      * @return the y-coordinate centroid
      */
     public double getY() {
-        return sumY / count;
+        // normalization required due to floating point precision errors
+        return GeoUtils.normalizeLat(sumY / sumWeight);
+    }
+
+    /**
+     * @return the sum of all the weighted coordinates summed in the calculator
+     */
+    public double sumWeight() {
+        return sumWeight;
     }
 
+    /**
+     * @return the highest dimensional shape type summed in the calculator
+     */
     public DimensionalShapeType getDimensionalShapeType() {
         return dimensionalShapeType;
     }
@@ -121,8 +139,7 @@ private CentroidCalculatorVisitor(CentroidCalculator calculator) {
 
         @Override
         public Void visit(Circle circle) {
-            calculator.addCoordinate(circle.getX(), circle.getY());
-            return null;
+            throw new IllegalArgumentException("invalid shape type found [Circle] while calculating centroid");
         }
 
         @Override
@@ -135,17 +152,47 @@ public Void visit(GeometryCollection<?> collection) {
 
         @Override
         public Void visit(Line line) {
-            for (int i = 0; i < line.length(); i++) {
-                calculator.addCoordinate(line.getX(i), line.getY(i));
+            // a line's centroid is calculated by summing the center of each
+            // line segment weighted by the line segment's length in degrees
+            for (int i = 0; i < line.length() - 1; i++) {
+                double diffX = line.getX(i) - line.getX(i + 1);
+                double diffY = line.getY(i) - line.getY(i + 1);
+                double x = (line.getX(i) + line.getX(i + 1)) / 2;
+                double y = (line.getY(i) + line.getY(i + 1)) / 2;
+                calculator.addCoordinate(x, y, Math.sqrt(diffX * diffX + diffY * diffY));
             }
             return null;
         }
-
         @Override
         public Void visit(LinearRing ring) {
+            throw new IllegalArgumentException("invalid shape type found [LinearRing] while calculating centroid");
+        }
+
+        private Void visit(LinearRing ring, boolean isHole) {
+            // implementation of calculation defined in
+            // https://www.seas.upenn.edu/~sys502/extra_materials/Polygon%20Area%20and%20Centroid.pdf
+            //
+            // centroid of a ring is a weighted coordinate based on the ring's area.
+            // the sign of the area is positive for the outer-shell of a polygon and negative for the holes
+
+            int sign = isHole ? -1 : 1;
+            double totalRingArea = 0.0;
             for (int i = 0; i < ring.length() - 1; i++) {
-                calculator.addCoordinate(ring.getX(i), ring.getY(i));
+                totalRingArea += (ring.getX(i) * ring.getY(i + 1)) - (ring.getX(i + 1) * ring.getY(i));
             }
+            totalRingArea = totalRingArea / 2;
+
+            double sumX = 0.0;
+            double sumY = 0.0;
+            for (int i = 0; i < ring.length() - 1; i++) {
+                double twiceArea = (ring.getX(i) * ring.getY(i + 1)) - (ring.getX(i + 1) * ring.getY(i));
+                sumX += twiceArea * (ring.getX(i) + ring.getX(i + 1));
+                sumY += twiceArea * (ring.getY(i) + ring.getY(i + 1));
+            }
+            double cX = sumX / (6 * totalRingArea);
+            double cY = sumY / (6 * totalRingArea);
+            calculator.addCoordinate(cX, cY, sign * Math.abs(totalRingArea));
+
             return null;
         }
 
@@ -175,22 +222,26 @@ public Void visit(MultiPolygon multiPolygon) {
 
         @Override
         public Void visit(Point point) {
-            calculator.addCoordinate(point.getX(), point.getY());
+            calculator.addCoordinate(point.getX(), point.getY(), 1.0);
             return null;
         }
 
         @Override
         public Void visit(Polygon polygon) {
-            // TODO: incorporate holes into centroid calculation
-            return visit(polygon.getPolygon());
+            visit(polygon.getPolygon(), false);
+            for (int i = 0; i < polygon.getNumberOfHoles(); i++) {
+                visit(polygon.getHole(i), true);
+            }
+            return null;
         }
 
         @Override
         public Void visit(Rectangle rectangle) {
-            calculator.addCoordinate(rectangle.getMinX(), rectangle.getMinY());
-            calculator.addCoordinate(rectangle.getMinX(), rectangle.getMaxY());
-            calculator.addCoordinate(rectangle.getMaxX(), rectangle.getMinY());
-            calculator.addCoordinate(rectangle.getMaxX(), rectangle.getMaxY());
+            double sumX = rectangle.getMaxX() + rectangle.getMinX();
+            double sumY = rectangle.getMaxY() + rectangle.getMinY();
+            double diffX = rectangle.getMaxX() - rectangle.getMinX();
+            double diffY = rectangle.getMaxY() - rectangle.getMinY();
+            calculator.addCoordinate(sumX / 2, sumY / 2, Math.abs(diffX * diffY));
             return null;
         }
     }

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

@@ -53,9 +53,9 @@ public enum DimensionalShapeType {
     GEOMETRYCOLLECTION_LINES,      // highest-dimensional shapes are Lines
     GEOMETRYCOLLECTION_POLYGONS;   // highest-dimensional shapes are Polygons
 
-    private static DimensionalShapeType[] values = values();
+    public static Comparator<DimensionalShapeType> COMPARATOR = Comparator.comparingInt(DimensionalShapeType::centroidDimension);
 
-    private static Comparator<DimensionalShapeType> COMPARATOR = Comparator.comparingInt(DimensionalShapeType::centroidDimension);
+    private static DimensionalShapeType[] values = values();
 
     public static DimensionalShapeType max(DimensionalShapeType s1, DimensionalShapeType s2) {
         if (s1 == null) {
@@ -66,12 +66,16 @@ public static DimensionalShapeType max(DimensionalShapeType s1, DimensionalShape
         return COMPARATOR.compare(s1, s2) >= 0 ? s1 : s2;
     }
 
+    public static DimensionalShapeType fromOrdinalByte(byte ordinal) {
+        return values[Byte.toUnsignedInt(ordinal)];
+    }
+
     public void writeTo(ByteBuffersDataOutput out) {
         out.writeByte((byte) ordinal());
     }
 
     public static DimensionalShapeType readFrom(ByteArrayDataInput in) {
-        return values[Byte.toUnsignedInt(in.readByte())];
+        return fromOrdinalByte(in.readByte());
     }
 
     public static DimensionalShapeType forGeometry(Geometry geometry) {
@@ -80,8 +84,7 @@ public static DimensionalShapeType forGeometry(Geometry geometry) {
 
             @Override
             public DimensionalShapeType visit(Circle circle) {
-                st = DimensionalShapeType.max(st, DimensionalShapeType.POLYGON);
-                return st;
+                throw new IllegalArgumentException("invalid shape type found [Circle] while computing dimensional shape type");
             }
 
             @Override

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

@@ -31,10 +31,26 @@
  *
  * This class supports checking bounding box
  * relations against the serialized triangle tree.
+ *
+ * -----------------------------------------
+ * |   The binary format of the tree       |
+ * -----------------------------------------
+ * -----------------------------------------  --
+ * |    centroid-x-coord (4 bytes)         |    |
+ * -----------------------------------------    |
+ * |    centroid-y-coord (4 bytes)         |    |
+ * -----------------------------------------    |
+ * |    DimensionalShapeType (1 byte)      |    | Centroid-related header
+ * -----------------------------------------    |
+ * |  Sum of weights (VLong 1-8 bytes)     |    |
+ * -----------------------------------------  --
+ * |         Extent (var-encoding)         |
+ * -----------------------------------------
+ * |         Triangle Tree                 |
+ * -----------------------------------------
+ * -----------------------------------------
  */
 public class TriangleTreeReader {
-    private static final int CENTROID_HEADER_SIZE_IN_BYTES = 9;
-
     private final ByteArrayDataInput input;
     private final CoordinateEncoder coordinateEncoder;
     private final Rectangle2D rectangle2D;
@@ -58,8 +74,7 @@ public void reset(BytesRef bytesRef) throws IOException {
      */
     public Extent getExtent() {
         if (treeOffset == 0) {
-            // TODO: Compress serialization of extent
-            input.setPosition(CENTROID_HEADER_SIZE_IN_BYTES);
+            getSumCentroidWeight(); // skip CENTROID_HEADER + var-long sum-weight
             Extent.readFromCompressed(input, extent);
             treeOffset = input.getPosition();
         } else {
@@ -89,6 +104,11 @@ public DimensionalShapeType getDimensionalShapeType() {
         return DimensionalShapeType.readFrom(input);
     }
 
+    public double getSumCentroidWeight() {
+        input.setPosition(9);
+        return Double.longBitsToDouble(input.readVLong());
+    }
+
     /**
      * Compute the relation with the provided bounding box. If the result is CELL_INSIDE_QUERY
      * then the bounding box is within the shape.

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

@@ -50,6 +50,7 @@ public void writeTo(ByteBuffersDataOutput out) throws IOException {
         out.writeInt(coordinateEncoder.encodeX(centroidCalculator.getX()));
         out.writeInt(coordinateEncoder.encodeY(centroidCalculator.getY()));
         centroidCalculator.getDimensionalShapeType().writeTo(out);
+        out.writeVLong(Double.doubleToLongBits(centroidCalculator.sumWeight()));
         extent.writeCompressed(out);
         node.writeTo(out);
     }

server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java

@@ -123,6 +123,11 @@ public DimensionalShapeType dimensionalShapeType() {
             return DimensionalShapeType.POINT;
         }
 
+        @Override
+        public double weight() {
+            return 1.0;
+        }
+
         @Override
         public double lat() {
             return geoPoint.lat();
@@ -173,6 +178,11 @@ public DimensionalShapeType dimensionalShapeType() {
             return reader.getDimensionalShapeType();
         }
 
+        @Override
+        public double weight() {
+            return reader.getSumCentroidWeight();
+        }
+
         @Override
         public double lat() {
             return reader.getCentroidY();
@@ -229,6 +239,7 @@ public interface GeoValue {
         BoundingBox boundingBox();
         GeoRelation relate(Rectangle rectangle);
         DimensionalShapeType dimensionalShapeType();
+        double weight();
     }
 
     public static class BoundingBox {