Fix geo_shape centroid calculation for downgraded shapes

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

@@ -31,29 +31,32 @@
 import org.elasticsearch.geometry.Point;
 import org.elasticsearch.geometry.Polygon;
 import org.elasticsearch.geometry.Rectangle;
+import org.elasticsearch.search.aggregations.metrics.CompensatedSum;
+
+import static org.elasticsearch.common.geo.DimensionalShapeType.LINE;
+import static org.elasticsearch.common.geo.DimensionalShapeType.POINT;
+import static org.elasticsearch.common.geo.DimensionalShapeType.POLYGON;
 
 /**
  * This class keeps a running Kahan-sum of coordinates
  * that are to be averaged in {@link TriangleTreeWriter} for use
  * as the centroid of a shape.
  */
 public class CentroidCalculator {
-    private double compX;
-    private double compY;
-    private double sumX;
-    private double sumY;
-    private double sumWeight;
+    CompensatedSum compSumX;
+    CompensatedSum compSumY;
+    CompensatedSum compSumWeight;
+    private CentroidCalculatorVisitor visitor;
     private DimensionalShapeType dimensionalShapeType;
 
     public CentroidCalculator(Geometry geometry) {
-        this.sumX = 0.0;
-        this.compX = 0.0;
-        this.sumY = 0.0;
-        this.compY = 0.0;
-        this.sumWeight = 0.0;
-        CentroidCalculatorVisitor visitor = new CentroidCalculatorVisitor(this);
+        this.compSumX = new CompensatedSum(0, 0);
+        this.compSumY = new CompensatedSum(0, 0);
+        this.compSumWeight = new CompensatedSum(0, 0);
+        this.dimensionalShapeType = null;
+        this.visitor = new CentroidCalculatorVisitor(this);
         geometry.visit(visitor);
-        this.dimensionalShapeType = DimensionalShapeType.forGeometry(geometry);
+        this.dimensionalShapeType = visitor.calculator.dimensionalShapeType;
     }
 
     /**
@@ -63,18 +66,19 @@ public CentroidCalculator(Geometry geometry) {
      * @param y the y-coordinate of the point
      * @param weight the associated weight of the coordinate
      */
-    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 = weight * y - compY;
-        double newSumY = sumY + correctedY;
-        compY = (newSumY - sumY) - correctedY;
-        sumY = newSumY;
-
-        sumWeight += weight;
+    private void addCoordinate(double x, double y, double weight, DimensionalShapeType dimensionalShapeType) {
+        if (this.dimensionalShapeType == null || this.dimensionalShapeType == dimensionalShapeType) {
+            compSumX.add(x * weight);
+            compSumY.add(y * weight);
+            compSumWeight.add(weight);
+            this.dimensionalShapeType = dimensionalShapeType;
+        } else if (dimensionalShapeType.compareTo(this.dimensionalShapeType) > 0) {
+            // reset counters
+            compSumX.reset(x * weight, 0);
+            compSumY.reset(y * weight, 0);
+            compSumWeight.reset(weight, 0);
+            this.dimensionalShapeType = dimensionalShapeType;
+        }
     }
 
     /**
@@ -86,16 +90,17 @@ private void addCoordinate(double x, double y, double weight) {
      * @param otherCalculator the other centroid calculator to add from
      */
     public void addFrom(CentroidCalculator otherCalculator) {
-        int compared = DimensionalShapeType.COMPARATOR.compare(dimensionalShapeType, otherCalculator.dimensionalShapeType);
+        int compared = dimensionalShapeType.compareTo(otherCalculator.dimensionalShapeType);
         if (compared < 0) {
-            sumWeight = otherCalculator.sumWeight;
             dimensionalShapeType = otherCalculator.dimensionalShapeType;
-            sumX = otherCalculator.sumX;
-            sumY = otherCalculator.sumY;
-            compX = otherCalculator.compX;
-            compY = otherCalculator.compY;
+            this.compSumX = otherCalculator.compSumX;
+            this.compSumY = otherCalculator.compSumY;
+            this.compSumWeight = otherCalculator.compSumWeight;
+
         } else if (compared == 0) {
-            addCoordinate(otherCalculator.sumX, otherCalculator.sumY, otherCalculator.sumWeight);
+            this.compSumX.add(otherCalculator.compSumX.value());
+            this.compSumY.add(otherCalculator.compSumY.value());
+            this.compSumWeight.add(otherCalculator.compSumWeight.value());
         } // else (compared > 0) do not modify centroid calculation since otherCalculator is of lower dimension than this calculator
     }
 
@@ -104,22 +109,22 @@ public void addFrom(CentroidCalculator otherCalculator) {
      */
     public double getX() {
         // normalization required due to floating point precision errors
-        return GeoUtils.normalizeLon(sumX / sumWeight);
+        return GeoUtils.normalizeLon(compSumX.value() / compSumWeight.value());
     }
 
     /**
      * @return the y-coordinate centroid
      */
     public double getY() {
         // normalization required due to floating point precision errors
-        return GeoUtils.normalizeLat(sumY / sumWeight);
+        return GeoUtils.normalizeLat(compSumY.value() / compSumWeight.value());
     }
 
     /**
      * @return the sum of all the weighted coordinates summed in the calculator
      */
     public double sumWeight() {
-        return sumWeight;
+        return compSumWeight.value();
     }
 
     /**
@@ -152,62 +157,34 @@ public Void visit(GeometryCollection<?> collection) {
 
         @Override
         public Void visit(Line line) {
-            // 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));
+            if (calculator.dimensionalShapeType != POLYGON) {
+                visitLine(line.length(), line::getX, line::getY);
             }
             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++) {
-                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;
-        }
 
         @Override
         public Void visit(MultiLine multiLine) {
-            for (Line line : multiLine) {
-                visit(line);
+            if (calculator.getDimensionalShapeType() != POLYGON) {
+                for (Line line : multiLine) {
+                    visit(line);
+                }
             }
             return null;
         }
 
         @Override
         public Void visit(MultiPoint multiPoint) {
-            for (Point point : multiPoint) {
-                visit(point);
+            if (calculator.getDimensionalShapeType() == null || calculator.getDimensionalShapeType() == POINT) {
+                for (Point point : multiPoint) {
+                    visit(point);
+                }
             }
             return null;
         }
@@ -222,16 +199,39 @@ public Void visit(MultiPolygon multiPolygon) {
 
         @Override
         public Void visit(Point point) {
-            calculator.addCoordinate(point.getX(), point.getY(), 1.0);
+            if (calculator.getDimensionalShapeType() == null || calculator.getDimensionalShapeType() == POINT) {
+                visitPoint(point.getX(), point.getY());
+            }
             return null;
         }
 
         @Override
         public Void visit(Polygon polygon) {
-            visit(polygon.getPolygon(), false);
+            // check area of polygon
+
+            double[] centroidX = new double[1 + polygon.getNumberOfHoles()];
+            double[] centroidY = new double[1 + polygon.getNumberOfHoles()];
+            double[] weight = new double[1 + polygon.getNumberOfHoles()];
+            visitLinearRing(polygon.getPolygon().length(), polygon.getPolygon()::getX, polygon.getPolygon()::getY, false,
+                centroidX, centroidY, weight, 0);
             for (int i = 0; i < polygon.getNumberOfHoles(); i++) {
-                visit(polygon.getHole(i), true);
+                visitLinearRing(polygon.getHole(i).length(), polygon.getHole(i)::getX, polygon.getHole(i)::getY, true,
+                    centroidX, centroidY, weight, i + 1);
+            }
+
+            double sumWeight = 0;
+            for (double w : weight) {
+                sumWeight += w;
+            }
+
+            if (sumWeight == 0 && calculator.dimensionalShapeType != POLYGON) {
+                visitLine(polygon.getPolygon().length(), polygon.getPolygon()::getX, polygon.getPolygon()::getY);
+            } else {
+                for (int i = 0; i < 1 + polygon.getNumberOfHoles(); i++) {
+                    calculator.addCoordinate(centroidX[i], centroidY[i], weight[i], POLYGON);
+                }
             }
+
             return null;
         }
 
@@ -241,9 +241,73 @@ public Void visit(Rectangle rectangle) {
             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));
+            if (diffX != 0 && diffY != 0) {
+                calculator.addCoordinate(sumX / 2, sumY / 2, Math.abs(diffX * diffY), POLYGON);
+            } else if (diffX != 0) {
+                calculator.addCoordinate(sumX / 2, rectangle.getMinY(), diffX, LINE);
+            } else if (diffY != 0) {
+                calculator.addCoordinate(rectangle.getMinX(), sumY / 2, diffY, LINE);
+            } else {
+                visitPoint(rectangle.getMinX(), rectangle.getMinY());
+            }
             return null;
         }
+
+
+        private void visitPoint(double x, double y) {
+            calculator.addCoordinate(x, y, 1.0, POINT);
+        }
+
+        private void visitLine(int length, CoordinateSupplier x, CoordinateSupplier y) {
+            // check line has length
+            double originDiffX = x.get(0) - x.get(1);
+            double originDiffY = y.get(0) - y.get(1);
+            if (originDiffX != 0 || originDiffY != 0) {
+                // 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 < length - 1; i++) {
+                    double diffX = x.get(i) - x.get(i + 1);
+                    double diffY = y.get(i) - y.get(i + 1);
+                    double xAvg = (x.get(i) + x.get(i + 1)) / 2;
+                    double yAvg = (y.get(i) + y.get(i + 1)) / 2;
+                    double weight = Math.sqrt(diffX * diffX + diffY * diffY);
+                    calculator.addCoordinate(xAvg, yAvg, weight, LINE);
+                }
+            } else {
+                visitPoint(x.get(0), y.get(0));
+            }
+        }
+
+        private void visitLinearRing(int length, CoordinateSupplier x, CoordinateSupplier y, boolean isHole,
+                                       double[] centroidX, double[] centroidY, double[] weight, int idx) {
+            // 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 < length - 1; i++) {
+                totalRingArea += (x.get(i) * y.get(i + 1)) - (x.get(i + 1) * y.get(i));
+            }
+            totalRingArea = totalRingArea / 2;
+
+            double sumX = 0.0;
+            double sumY = 0.0;
+            for (int i = 0; i < length - 1; i++) {
+                double twiceArea = (x.get(i) * y.get(i + 1)) - (x.get(i + 1) * y.get(i));
+                sumX += twiceArea * (x.get(i) + x.get(i + 1));
+                sumY += twiceArea * (y.get(i) + y.get(i + 1));
+            }
+            centroidX[idx] = sumX / (6 * totalRingArea);
+            centroidY[idx] = sumY / (6 * totalRingArea);
+            weight[idx] = sign * Math.abs(totalRingArea);
+        }
     }
 
+    @FunctionalInterface
+    private interface CoordinateSupplier {
+        double get(int idx);
+    }
 }